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
|
package server
import (
"tunnel/pkg/config"
"strings"
"bytes"
"sync"
"time"
"fmt"
"log"
"net"
"io"
)
type handler func (*Server, []string, io.Writer)
var handlers = map[string]handler{}
type Server struct {
listen net.Listener
since time.Time
wg sync.WaitGroup
m sync.Mutex
done bool
}
func setHandler(cmd string, h handler) {
if _, ok := handlers[cmd]; ok {
panic(fmt.Sprintf("handler %s already registered", cmd))
}
handlers[cmd] = h
}
func (s *Server) isDone() bool {
s.m.Lock()
defer s.m.Unlock()
return s.done
}
func New() (*Server, error) {
listen, err := net.Listen(config.SockType, config.SockPath)
if err != nil {
return nil, err
}
s := &Server{
listen: listen,
since: time.Now(),
}
return s, nil
}
func (s *Server) Run() {
for {
conn, err := s.listen.Accept()
if err != nil {
if s.isDone() {
break
}
log.Print(err)
continue
}
log.Print("new client")
s.wg.Add(1)
go s.handle(conn)
}
s.wg.Wait()
}
func (s *Server) Stop() {
s.m.Lock()
s.done = true
s.m.Unlock()
s.listen.Close()
}
func (s *Server) handle(c net.Conn) {
defer c.Close()
defer s.wg.Done()
buf := make([]byte, config.BufSize)
nr, er := c.Read(buf)
if er != nil {
if er != io.EOF {
log.Print("handle: ", er)
}
return
}
data := string(buf[:nr])
args := strings.Split(data, " ")
out := bytes.NewBuffer(nil)
if h, ok := handlers[args[0]]; ok {
log.Printf("handle: run: %s", data)
h(s, args, out)
} else {
fmt.Fprint(out, "unknown command")
}
if out.Len() == 0 {
out.Write([]byte("\n"))
}
_, ew := c.Write(out.Bytes())
if ew != nil {
log.Print("handle: ", ew)
return
}
}
|
