-
Notifications
You must be signed in to change notification settings - Fork 32
/
tcp.go
103 lines (84 loc) · 1.97 KB
/
tcp.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
package tcp
import (
"fmt"
"log"
"net"
"github.com/packetzoom/logzoom/buffer"
"github.com/packetzoom/logzoom/output"
"gopkg.in/yaml.v2"
)
const (
recvBuffer = 100
)
type Config struct {
Host string `yaml:"host"`
}
type TCPServer struct {
host string
b buffer.Sender
term chan bool
}
func init() {
output.Register("tcp", &TCPServer{
term: make(chan bool, 1),
})
}
// lumberConn handles an incoming connection from a lumberjack client
func (s *TCPServer) accept(c net.Conn) {
defer func() {
s.b.DelSubscriber(c.RemoteAddr().String())
log.Printf("[%s] closing tcp connection", c.RemoteAddr().String())
c.Close()
}()
log.Printf("[%s] accepting tcp connection", c.RemoteAddr().String())
// Add the client as a subscriber
r := make(chan *buffer.Event, recvBuffer)
s.b.AddSubscriber(c.RemoteAddr().String(), r)
for {
select {
case ev := <-r:
_, err := c.Write([]byte(fmt.Sprintf("%s %s\n", ev.Source, *ev.Text)))
if err != nil {
log.Printf("[%s] error sending event to tcp connection: %v", c.RemoteAddr().String(), err)
return
}
}
}
}
func (s *TCPServer) Init(config yaml.MapSlice, b buffer.Sender) error {
var tcpConfig *Config
// go-yaml doesn't have a great way to partially unmarshal YAML data
// See https://github.com/go-yaml/yaml/issues/13
yamlConfig, _ := yaml.Marshal(config)
if err := yaml.Unmarshal(yamlConfig, &tcpConfig); err != nil {
return fmt.Errorf("Error parsing tcp config: %v", err)
}
s.host = tcpConfig.Host
s.b = b
return nil
}
func (s *TCPServer) Start() error {
ln, err := net.Listen("tcp", s.host)
if err != nil {
return fmt.Errorf("TCPServer: listener failed: %v", err)
}
for {
select {
case <-s.term:
log.Println("TCPServer received term signal")
return nil
default:
conn, err := ln.Accept()
if err != nil {
log.Println("Error accepting tcp connection: %v", err)
continue
}
go s.accept(conn)
}
}
return nil
}
func (s *TCPServer) Stop() error {
s.term <- true
return nil
}