github.com/imgk/caddy-trojan@v0.0.0-20221206043256-2631719e16c8/listener/listener.go (about)

     1  package listener
     2  
     3  import (
     4  	"errors"
     5  	"fmt"
     6  	"io"
     7  	"net"
     8  	"os"
     9  
    10  	"github.com/caddyserver/caddy/v2"
    11  	"github.com/caddyserver/caddy/v2/caddyconfig/caddyfile"
    12  
    13  	"go.uber.org/zap"
    14  
    15  	"github.com/imgk/caddy-trojan/app"
    16  	"github.com/imgk/caddy-trojan/trojan"
    17  	"github.com/imgk/caddy-trojan/utils"
    18  )
    19  
    20  func init() {
    21  	caddy.RegisterModule(ListenerWrapper{})
    22  }
    23  
    24  // ListenerWrapper implements an TLS wrapper that it accept connections
    25  // from clients and check the connection with pre-defined password
    26  // and aead cipher defined by go-shadowsocks2, and return a normal page if
    27  // failed.
    28  type ListenerWrapper struct {
    29  	// Upstream is ...
    30  	Upstream app.Upstream `json:"-,omitempty"`
    31  	// Proxy is ...
    32  	Proxy app.Proxy `json:"-,omitempty"`
    33  	// Logger is ...
    34  	Logger *zap.Logger `json:"-,omitempty"`
    35  	// Verbose is ...
    36  	Verbose bool `json:"verbose,omitempty"`
    37  }
    38  
    39  // CaddyModule returns the Caddy module information.
    40  func (ListenerWrapper) CaddyModule() caddy.ModuleInfo {
    41  	return caddy.ModuleInfo{
    42  		ID:  "caddy.listeners.trojan",
    43  		New: func() caddy.Module { return new(ListenerWrapper) },
    44  	}
    45  }
    46  
    47  // Provision implements caddy.Provisioner.
    48  func (m *ListenerWrapper) Provision(ctx caddy.Context) error {
    49  	m.Logger = ctx.Logger(m)
    50  	if !ctx.AppIsConfigured(app.CaddyAppID) {
    51  		return errors.New("listener: trojan is not configured")
    52  	}
    53  	mod, err := ctx.App(app.CaddyAppID)
    54  	if err != nil {
    55  		return err
    56  	}
    57  	app := mod.(*app.App)
    58  	m.Upstream = app.Upstream()
    59  	m.Proxy = app.Proxy()
    60  	return nil
    61  }
    62  
    63  // WrapListener implements caddy.ListenWrapper
    64  func (m *ListenerWrapper) WrapListener(l net.Listener) net.Listener {
    65  	ln := NewListener(l, m.Upstream, m.Proxy, m.Logger)
    66  	ln.Verbose = m.Verbose
    67  	go ln.loop()
    68  	return ln
    69  }
    70  
    71  // UnmarshalCaddyfile unmarshals Caddyfile tokens into h.
    72  func (*ListenerWrapper) UnmarshalCaddyfile(d *caddyfile.Dispenser) error {
    73  	return nil
    74  }
    75  
    76  // Interface guards
    77  var (
    78  	_ caddy.Provisioner     = (*ListenerWrapper)(nil)
    79  	_ caddy.ListenerWrapper = (*ListenerWrapper)(nil)
    80  	_ caddyfile.Unmarshaler = (*ListenerWrapper)(nil)
    81  )
    82  
    83  // Listener is ...
    84  type Listener struct {
    85  	Verbose bool
    86  
    87  	// Listener is ...
    88  	net.Listener
    89  	// Upstream is ...
    90  	Upstream app.Upstream
    91  	// Proxy is ...
    92  	Proxy app.Proxy
    93  	// Logger is ...
    94  	Logger *zap.Logger
    95  
    96  	// return *rawConn
    97  	conns chan net.Conn
    98  	// close channel
    99  	closed chan struct{}
   100  }
   101  
   102  // NewListener is ...
   103  func NewListener(ln net.Listener, up app.Upstream, px app.Proxy, logger *zap.Logger) *Listener {
   104  	l := &Listener{
   105  		Listener: ln,
   106  		Upstream: up,
   107  		Proxy:    px,
   108  		Logger:   logger,
   109  		conns:    make(chan net.Conn, 8),
   110  		closed:   make(chan struct{}),
   111  	}
   112  	return l
   113  }
   114  
   115  // Accept is ...
   116  func (l *Listener) Accept() (net.Conn, error) {
   117  	select {
   118  	case <-l.closed:
   119  		return nil, os.ErrClosed
   120  	case c := <-l.conns:
   121  		return c, nil
   122  	}
   123  }
   124  
   125  // Close is ...
   126  func (l *Listener) Close() error {
   127  	select {
   128  	case <-l.closed:
   129  		return nil
   130  	default:
   131  		close(l.closed)
   132  	}
   133  	return nil
   134  }
   135  
   136  // loop is ...
   137  func (l *Listener) loop() {
   138  	for {
   139  		conn, err := l.Listener.Accept()
   140  		if err != nil {
   141  			select {
   142  			case <-l.closed:
   143  				return
   144  			default:
   145  				l.Logger.Error(fmt.Sprintf("accept net.Conn error: %v", err))
   146  			}
   147  			continue
   148  		}
   149  
   150  		go func(c net.Conn, lg *zap.Logger, up app.Upstream) {
   151  			b := make([]byte, trojan.HeaderLen+2)
   152  			for n := 0; n < trojan.HeaderLen+2; n += 1 {
   153  				nr, err := c.Read(b[n : n+1])
   154  				if err != nil {
   155  					if errors.Is(err, io.EOF) {
   156  						lg.Error(fmt.Sprintf("read prefix error: read tcp %v -> %v: read: %v", c.RemoteAddr(), c.LocalAddr(), err))
   157  					} else {
   158  						lg.Error(fmt.Sprintf("read prefix error, not io, rewind and let normal caddy deal with it: %v", err))
   159  						l.conns <- utils.RewindConn(c, b[:n+1])
   160  						return
   161  					}
   162  					c.Close()
   163  					return
   164  				}
   165  				if nr == 0 {
   166  					continue
   167  				}
   168  				// mimic nginx
   169  				if b[n] == 0x0a && n < trojan.HeaderLen+1 {
   170  					select {
   171  					case <-l.closed:
   172  						c.Close()
   173  					default:
   174  						l.conns <- utils.RewindConn(c, b[:n+1])
   175  					}
   176  					return
   177  				}
   178  			}
   179  
   180  			// check the net.Conn
   181  			if ok := up.Validate(utils.ByteSliceToString(b[:trojan.HeaderLen])); !ok {
   182  				select {
   183  				case <-l.closed:
   184  					c.Close()
   185  				default:
   186  					l.conns <- utils.RewindConn(c, b)
   187  				}
   188  				return
   189  			}
   190  			defer c.Close()
   191  			if l.Verbose {
   192  				lg.Info(fmt.Sprintf("handle trojan net.Conn from %v", c.RemoteAddr()))
   193  			}
   194  
   195  			nr, nw, err := l.Proxy.Handle(io.Reader(c), io.Writer(c))
   196  			if err != nil {
   197  				lg.Error(fmt.Sprintf("handle net.Conn error: %v", err))
   198  			}
   199  			up.Consume(utils.ByteSliceToString(b[:trojan.HeaderLen]), nr, nw)
   200  		}(conn, l.Logger, l.Upstream)
   201  	}
   202  }