-
Notifications
You must be signed in to change notification settings - Fork 12
/
tcptee.go
108 lines (96 loc) · 2.01 KB
/
tcptee.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
// Copyright 2016 Chao Wang <hit9@icloud.com>
// Package main implements a tcp tee.
// Usage: ./tcptee -bind :8000 -backends :2015,:2016,:2017
package main
import (
"flag"
"io"
"log"
"net"
"strings"
)
// Backends is the backend connections.
type Backends []net.Conn
// Read implements the io.Reader.
func (b Backends) Read(p []byte) (n int, err error) {
for _, conn := range b {
n, err = conn.Read(p)
if err != nil {
return n, err
}
}
return
}
// Write implements the io.Writer.
func (b Backends) Write(p []byte) (n int, err error) {
for _, conn := range b {
n, err = conn.Write(p)
if err != nil {
return n, err
}
}
return
}
// Tee is the tee handle.
type Tee struct {
ln net.Listener
laddr string // server addr
addrs []string // backend addrs
}
// New creates a new Tee.
func New(laddr string, addrs []string) *Tee {
return &Tee{laddr: laddr, addrs: addrs}
}
// Listen the tee.
func (t *Tee) Listen() (err error) {
t.ln, err = net.Listen("tcp", t.laddr)
if err != nil {
return err
}
log.Printf("tee is listening on %s\n", t.laddr)
return nil
}
// Serve the tee.
func (t *Tee) Serve() error {
for {
conn, err := t.ln.Accept()
if err != nil {
return err
}
go t.handle(conn)
}
}
// ListenAndServe is the Listen followed by Serve.
func (t *Tee) ListenAndServe() (err error) {
if err = t.Listen(); err != nil {
return
}
return t.Serve()
}
// Handle the connection.
func (t *Tee) handle(conn net.Conn) {
// Connect to backends
var backends []net.Conn
for _, addr := range t.addrs {
c, err := net.Dial("tcp", addr)
if err != nil {
log.Println(err)
}
backends = append(backends, c)
}
var err error
_, err = io.Copy(Backends(backends), conn)
if err != nil {
log.Println(err)
}
}
func main() {
bind := flag.String("bind", ":8000", "address to bind")
backends := flag.String("backends", "", "backends split by comma")
flag.Parse()
if *backends == "" {
log.Fatal("No backends")
}
tr := New(*bind, strings.Split(*backends, ","))
log.Fatal(tr.ListenAndServe())
}