/
main.go
145 lines (125 loc) · 3.35 KB
/
main.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
package main
import (
"bufio"
"flag"
"fmt"
"io"
"log"
"mirgit.dcs.gla.ac.uk/JamesMcMinn/twitter"
"net"
"os"
"runtime"
"strconv"
)
const RECV_BUF_LEN = 1024 * 1024
const MAX_CHAN_LEN = 10000
const MODE_STREAM = 0
const MODE_FILE = 1
var (
consumerKey *string = flag.String("ck", "", "Consumer Key")
consumerSecret *string = flag.String("cs", "", "Consumer Secret")
ot *string = flag.String("ot", "", "OAuth Token")
osec *string = flag.String("os", "", "OAuthTokenSecret")
inputFile *string = flag.String("if", "", "Input File")
port *int = flag.Int("port", 8053, "Port to listen on. Default: 8053")
firehose chan twitter.Tweet = make(chan twitter.Tweet, MAX_CHAN_LEN)
aliveStreams map[chan *[]byte]bool = make(map[chan *[]byte]bool)
mode int = -1
fileFormat int = -1
)
func main() {
runtime.GOMAXPROCS(runtime.NumCPU() * 2)
flag.Parse()
if *inputFile != "" {
mode = MODE_FILE
} else if *consumerKey != "" || *consumerSecret != "" || *ot != "" || *osec != "" {
if *consumerKey == "" || *consumerSecret == "" || *ot == "" || *osec == "" {
fmt.Println("Must specify all of -ck, -cs, -ot and -os. See -help for details.")
return
}
mode = MODE_STREAM
} else {
fmt.Println("Must specify either Twitter OAuth details or file location and format. See -help for details.")
return
}
// Listen on whatever port was specified
ln, err := net.Listen("tcp", ":"+strconv.Itoa(*port))
if err != nil {
log.Fatal(err)
}
log.Println("Listening on port", *port)
if mode == MODE_STREAM {
// Open a connection the the firehose and fill output streams
go twitter.FillStream(firehose, *consumerKey, *consumerSecret, *ot, *osec)
go fillOutgoingStreams(aliveStreams)
}
for {
conn, err := ln.Accept()
if err != nil {
log.Println(err)
}
log.Println("New connection attempt from " + conn.RemoteAddr().String())
go handleConnection(conn)
}
}
// Reads a file into a channel
func readFileInto(into chan *[]byte) {
f, err := os.Open(*inputFile)
if err != nil {
log.Fatal(err)
}
bf := bufio.NewReaderSize(f, 20000)
for {
line, isPrefix, err := bf.ReadLine()
switch {
case err == io.EOF:
break
case err != nil:
log.Fatal(err)
case isPrefix:
log.Fatal("Error: Unexpected long line reading", f.Name())
}
// Check the connection is still active
if aliveStreams[into] != true {
break
}
t := twitter.JSONtoTweet(line)
j, err := twitter.TweetToJSON(t)
if err != nil {
log.Println(err)
}
j = append(j, []byte("\n")...)
into <- &j
}
}
func handleConnection(conn net.Conn) {
stream := make(chan *[]byte, MAX_CHAN_LEN)
aliveStreams[stream] = true
log.Println("Current Connections:", len(aliveStreams))
if mode == MODE_FILE {
go readFileInto(stream)
}
for {
t := <-stream
_, err := conn.Write(*t)
if err != nil {
log.Println("Closing connection: ", err.Error())
break
}
}
delete(aliveStreams, stream)
log.Println("Current Connections:", len(aliveStreams))
}
func fillOutgoingStreams(streams map[chan *[]byte]bool) {
for {
tweet := <-firehose
for r := range streams {
if len(r) == MAX_CHAN_LEN {
<-r
}
json, _ := twitter.TweetToJSON(tweet)
json = append(json, []byte("\n")...)
r <- &json
}
}
}