mirror of https://github.com/ginuerzh/gost
ginuerzh
6 years ago
19 changed files with 756 additions and 590 deletions
@ -0,0 +1,66 @@ |
|||||
|
package gost |
||||
|
|
||||
|
import "net" |
||||
|
|
||||
|
// tcpTransporter is a raw TCP transporter.
|
||||
|
type tcpTransporter struct{} |
||||
|
|
||||
|
// TCPTransporter creates a raw TCP client.
|
||||
|
func TCPTransporter() Transporter { |
||||
|
return &tcpTransporter{} |
||||
|
} |
||||
|
|
||||
|
func (tr *tcpTransporter) Dial(addr string, options ...DialOption) (net.Conn, error) { |
||||
|
opts := &DialOptions{} |
||||
|
for _, option := range options { |
||||
|
option(opts) |
||||
|
} |
||||
|
|
||||
|
timeout := opts.Timeout |
||||
|
if timeout <= 0 { |
||||
|
timeout = DialTimeout |
||||
|
} |
||||
|
if opts.Chain == nil { |
||||
|
return net.DialTimeout("tcp", addr, timeout) |
||||
|
} |
||||
|
return opts.Chain.Dial(addr) |
||||
|
} |
||||
|
|
||||
|
func (tr *tcpTransporter) Handshake(conn net.Conn, options ...HandshakeOption) (net.Conn, error) { |
||||
|
return conn, nil |
||||
|
} |
||||
|
|
||||
|
func (tr *tcpTransporter) Multiplex() bool { |
||||
|
return false |
||||
|
} |
||||
|
|
||||
|
type tcpListener struct { |
||||
|
net.Listener |
||||
|
} |
||||
|
|
||||
|
// TCPListener creates a Listener for TCP proxy server.
|
||||
|
func TCPListener(addr string) (Listener, error) { |
||||
|
laddr, err := net.ResolveTCPAddr("tcp", addr) |
||||
|
if err != nil { |
||||
|
return nil, err |
||||
|
} |
||||
|
ln, err := net.ListenTCP("tcp", laddr) |
||||
|
if err != nil { |
||||
|
return nil, err |
||||
|
} |
||||
|
return &tcpListener{Listener: tcpKeepAliveListener{ln}}, nil |
||||
|
} |
||||
|
|
||||
|
type tcpKeepAliveListener struct { |
||||
|
*net.TCPListener |
||||
|
} |
||||
|
|
||||
|
func (ln tcpKeepAliveListener) Accept() (c net.Conn, err error) { |
||||
|
tc, err := ln.AcceptTCP() |
||||
|
if err != nil { |
||||
|
return |
||||
|
} |
||||
|
tc.SetKeepAlive(true) |
||||
|
tc.SetKeepAlivePeriod(KeepAliveTime) |
||||
|
return tc, nil |
||||
|
} |
||||
@ -0,0 +1,357 @@ |
|||||
|
package gost |
||||
|
|
||||
|
import ( |
||||
|
"errors" |
||||
|
"net" |
||||
|
"sync" |
||||
|
"sync/atomic" |
||||
|
"time" |
||||
|
|
||||
|
"github.com/go-log/log" |
||||
|
) |
||||
|
|
||||
|
// udpTransporter is a raw UDP transporter.
|
||||
|
type udpTransporter struct{} |
||||
|
|
||||
|
// UDPTransporter creates a Transporter for UDP client.
|
||||
|
func UDPTransporter() Transporter { |
||||
|
return &udpTransporter{} |
||||
|
} |
||||
|
|
||||
|
func (tr *udpTransporter) Dial(addr string, options ...DialOption) (net.Conn, error) { |
||||
|
raddr, err := net.ResolveUDPAddr("udp", addr) |
||||
|
if err != nil { |
||||
|
return nil, err |
||||
|
} |
||||
|
|
||||
|
conn, err := net.ListenUDP("udp", nil) |
||||
|
if err != nil { |
||||
|
return nil, err |
||||
|
} |
||||
|
|
||||
|
return &udpClientConn{ |
||||
|
UDPConn: conn, |
||||
|
raddr: raddr, |
||||
|
}, nil |
||||
|
} |
||||
|
|
||||
|
func (tr *udpTransporter) Handshake(conn net.Conn, options ...HandshakeOption) (net.Conn, error) { |
||||
|
return conn, nil |
||||
|
} |
||||
|
|
||||
|
func (tr *udpTransporter) Multiplex() bool { |
||||
|
return false |
||||
|
} |
||||
|
|
||||
|
// UDPListenConfig is the config for UDP Listener.
|
||||
|
type UDPListenConfig struct { |
||||
|
TTL time.Duration // timeout per connection
|
||||
|
Backlog int // connection backlog
|
||||
|
QueueSize int // recv queue size per connection
|
||||
|
} |
||||
|
|
||||
|
type udpListener struct { |
||||
|
ln net.PacketConn |
||||
|
connChan chan net.Conn |
||||
|
errChan chan error |
||||
|
connMap udpConnMap |
||||
|
config *UDPListenConfig |
||||
|
} |
||||
|
|
||||
|
// UDPListener creates a Listener for UDP server.
|
||||
|
func UDPListener(addr string, cfg *UDPListenConfig) (Listener, error) { |
||||
|
laddr, err := net.ResolveUDPAddr("udp", addr) |
||||
|
if err != nil { |
||||
|
return nil, err |
||||
|
} |
||||
|
ln, err := net.ListenUDP("udp", laddr) |
||||
|
if err != nil { |
||||
|
return nil, err |
||||
|
} |
||||
|
|
||||
|
if cfg == nil { |
||||
|
cfg = &UDPListenConfig{} |
||||
|
} |
||||
|
|
||||
|
backlog := cfg.Backlog |
||||
|
if backlog <= 0 { |
||||
|
backlog = defaultBacklog |
||||
|
} |
||||
|
|
||||
|
l := &udpListener{ |
||||
|
ln: ln, |
||||
|
connChan: make(chan net.Conn, backlog), |
||||
|
errChan: make(chan error, 1), |
||||
|
config: cfg, |
||||
|
} |
||||
|
go l.listenLoop() |
||||
|
return l, nil |
||||
|
} |
||||
|
|
||||
|
func (l *udpListener) listenLoop() { |
||||
|
for { |
||||
|
b := make([]byte, mediumBufferSize) |
||||
|
n, raddr, err := l.ln.ReadFrom(b) |
||||
|
if err != nil { |
||||
|
log.Logf("[udp] peer -> %s : %s", l.Addr(), err) |
||||
|
l.Close() |
||||
|
l.errChan <- err |
||||
|
close(l.errChan) |
||||
|
return |
||||
|
} |
||||
|
|
||||
|
conn, ok := l.connMap.Get(raddr.String()) |
||||
|
if !ok { |
||||
|
conn = newUDPServerConn(l.ln, raddr, &udpServerConnConfig{ |
||||
|
ttl: l.config.TTL, |
||||
|
qsize: l.config.QueueSize, |
||||
|
onClose: func() { |
||||
|
l.connMap.Delete(raddr.String()) |
||||
|
log.Logf("[udp] %s closed (%d)", raddr, l.connMap.Size()) |
||||
|
}, |
||||
|
}) |
||||
|
|
||||
|
select { |
||||
|
case l.connChan <- conn: |
||||
|
l.connMap.Set(raddr.String(), conn) |
||||
|
log.Logf("[udp] %s -> %s (%d)", raddr, l.Addr(), l.connMap.Size()) |
||||
|
default: |
||||
|
conn.Close() |
||||
|
log.Logf("[udp] %s - %s: connection queue is full (%d)", raddr, l.Addr(), cap(l.connChan)) |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
select { |
||||
|
case conn.rChan <- b[:n]: |
||||
|
if Debug { |
||||
|
log.Logf("[udp] %s >>> %s : length %d", raddr, l.Addr(), n) |
||||
|
} |
||||
|
default: |
||||
|
log.Logf("[udp] %s -> %s : recv queue is full (%d)", raddr, l.Addr(), cap(conn.rChan)) |
||||
|
} |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
func (l *udpListener) Accept() (conn net.Conn, err error) { |
||||
|
var ok bool |
||||
|
select { |
||||
|
case conn = <-l.connChan: |
||||
|
case err, ok = <-l.errChan: |
||||
|
if !ok { |
||||
|
err = errors.New("accpet on closed listener") |
||||
|
} |
||||
|
} |
||||
|
return |
||||
|
} |
||||
|
|
||||
|
func (l *udpListener) Addr() net.Addr { |
||||
|
return l.ln.LocalAddr() |
||||
|
} |
||||
|
|
||||
|
func (l *udpListener) Close() error { |
||||
|
err := l.ln.Close() |
||||
|
l.connMap.Range(func(k interface{}, v *udpServerConn) bool { |
||||
|
v.Close() |
||||
|
return true |
||||
|
}) |
||||
|
|
||||
|
return err |
||||
|
} |
||||
|
|
||||
|
type udpConnMap struct { |
||||
|
m sync.Map |
||||
|
size int64 |
||||
|
} |
||||
|
|
||||
|
func (m *udpConnMap) Get(key interface{}) (conn *udpServerConn, ok bool) { |
||||
|
v, ok := m.m.Load(key) |
||||
|
if ok { |
||||
|
conn, ok = v.(*udpServerConn) |
||||
|
} |
||||
|
return |
||||
|
} |
||||
|
|
||||
|
func (m *udpConnMap) Set(key interface{}, conn *udpServerConn) { |
||||
|
m.m.Store(key, conn) |
||||
|
atomic.AddInt64(&m.size, 1) |
||||
|
} |
||||
|
|
||||
|
func (m *udpConnMap) Delete(key interface{}) { |
||||
|
m.m.Delete(key) |
||||
|
atomic.AddInt64(&m.size, -1) |
||||
|
} |
||||
|
|
||||
|
func (m *udpConnMap) Range(f func(key interface{}, value *udpServerConn) bool) { |
||||
|
m.m.Range(func(k, v interface{}) bool { |
||||
|
return f(k, v.(*udpServerConn)) |
||||
|
}) |
||||
|
} |
||||
|
|
||||
|
func (m *udpConnMap) Size() int64 { |
||||
|
return atomic.LoadInt64(&m.size) |
||||
|
} |
||||
|
|
||||
|
// udpServerConn is a server side connection for UDP client peer, it implements net.Conn and net.PacketConn.
|
||||
|
type udpServerConn struct { |
||||
|
conn net.PacketConn |
||||
|
raddr net.Addr |
||||
|
rChan chan []byte |
||||
|
closed chan struct{} |
||||
|
closeMutex sync.Mutex |
||||
|
nopChan chan int |
||||
|
config *udpServerConnConfig |
||||
|
} |
||||
|
|
||||
|
type udpServerConnConfig struct { |
||||
|
ttl time.Duration |
||||
|
qsize int |
||||
|
onClose func() |
||||
|
} |
||||
|
|
||||
|
func newUDPServerConn(conn net.PacketConn, raddr net.Addr, cfg *udpServerConnConfig) *udpServerConn { |
||||
|
if conn == nil || raddr == nil { |
||||
|
return nil |
||||
|
} |
||||
|
|
||||
|
if cfg == nil { |
||||
|
cfg = &udpServerConnConfig{} |
||||
|
} |
||||
|
qsize := cfg.qsize |
||||
|
if qsize <= 0 { |
||||
|
qsize = defaultQueueSize |
||||
|
} |
||||
|
c := &udpServerConn{ |
||||
|
conn: conn, |
||||
|
raddr: raddr, |
||||
|
rChan: make(chan []byte, qsize), |
||||
|
closed: make(chan struct{}), |
||||
|
nopChan: make(chan int), |
||||
|
config: cfg, |
||||
|
} |
||||
|
go c.ttlWait() |
||||
|
return c |
||||
|
} |
||||
|
|
||||
|
func (c *udpServerConn) Read(b []byte) (n int, err error) { |
||||
|
n, _, err = c.ReadFrom(b) |
||||
|
return |
||||
|
} |
||||
|
|
||||
|
func (c *udpServerConn) ReadFrom(b []byte) (n int, addr net.Addr, err error) { |
||||
|
select { |
||||
|
case bb := <-c.rChan: |
||||
|
n = copy(b, bb) |
||||
|
case <-c.closed: |
||||
|
err = errors.New("read from closed connection") |
||||
|
return |
||||
|
} |
||||
|
|
||||
|
select { |
||||
|
case c.nopChan <- n: |
||||
|
default: |
||||
|
} |
||||
|
|
||||
|
addr = c.raddr |
||||
|
|
||||
|
return |
||||
|
} |
||||
|
|
||||
|
func (c *udpServerConn) Write(b []byte) (n int, err error) { |
||||
|
return c.WriteTo(b, c.raddr) |
||||
|
} |
||||
|
|
||||
|
func (c *udpServerConn) WriteTo(b []byte, addr net.Addr) (n int, err error) { |
||||
|
n, err = c.conn.WriteTo(b, addr) |
||||
|
|
||||
|
if n > 0 { |
||||
|
if Debug { |
||||
|
log.Logf("[udp] %s <<< %s : length %d", addr, c.LocalAddr(), n) |
||||
|
} |
||||
|
|
||||
|
select { |
||||
|
case c.nopChan <- n: |
||||
|
default: |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
return |
||||
|
} |
||||
|
|
||||
|
func (c *udpServerConn) Close() error { |
||||
|
c.closeMutex.Lock() |
||||
|
defer c.closeMutex.Unlock() |
||||
|
|
||||
|
select { |
||||
|
case <-c.closed: |
||||
|
return errors.New("connection is closed") |
||||
|
default: |
||||
|
if c.config.onClose != nil { |
||||
|
c.config.onClose() |
||||
|
} |
||||
|
close(c.closed) |
||||
|
} |
||||
|
return nil |
||||
|
} |
||||
|
|
||||
|
func (c *udpServerConn) ttlWait() { |
||||
|
ttl := c.config.ttl |
||||
|
if ttl == 0 { |
||||
|
ttl = defaultTTL |
||||
|
} |
||||
|
timer := time.NewTimer(ttl) |
||||
|
defer timer.Stop() |
||||
|
|
||||
|
for { |
||||
|
select { |
||||
|
case <-c.nopChan: |
||||
|
if !timer.Stop() { |
||||
|
<-timer.C |
||||
|
} |
||||
|
timer.Reset(ttl) |
||||
|
case <-timer.C: |
||||
|
c.Close() |
||||
|
return |
||||
|
case <-c.closed: |
||||
|
return |
||||
|
} |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
func (c *udpServerConn) LocalAddr() net.Addr { |
||||
|
return c.conn.LocalAddr() |
||||
|
} |
||||
|
|
||||
|
func (c *udpServerConn) RemoteAddr() net.Addr { |
||||
|
return c.raddr |
||||
|
} |
||||
|
|
||||
|
func (c *udpServerConn) SetDeadline(t time.Time) error { |
||||
|
return c.conn.SetDeadline(t) |
||||
|
} |
||||
|
|
||||
|
func (c *udpServerConn) SetReadDeadline(t time.Time) error { |
||||
|
return c.conn.SetReadDeadline(t) |
||||
|
} |
||||
|
|
||||
|
func (c *udpServerConn) SetWriteDeadline(t time.Time) error { |
||||
|
return c.conn.SetWriteDeadline(t) |
||||
|
} |
||||
|
|
||||
|
type udpClientConn struct { |
||||
|
*net.UDPConn |
||||
|
raddr net.Addr |
||||
|
} |
||||
|
|
||||
|
func (c *udpClientConn) Write(b []byte) (int, error) { |
||||
|
if c.raddr != nil { |
||||
|
return c.WriteTo(b, c.raddr) |
||||
|
} |
||||
|
return c.UDPConn.Write(b) |
||||
|
} |
||||
|
|
||||
|
func (c *udpClientConn) RemoteAddr() net.Addr { |
||||
|
if c.raddr != nil { |
||||
|
return c.raddr |
||||
|
} |
||||
|
return c.UDPConn.RemoteAddr() |
||||
|
} |
||||
Loading…
Reference in new issue