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

     1  package handler
     2  
     3  import (
     4  	"errors"
     5  	"fmt"
     6  	"io"
     7  	"net/http"
     8  	"strings"
     9  
    10  	"github.com/caddyserver/caddy/v2"
    11  	"github.com/caddyserver/caddy/v2/caddyconfig/caddyfile"
    12  	"github.com/caddyserver/caddy/v2/caddyconfig/httpcaddyfile"
    13  	"github.com/caddyserver/caddy/v2/modules/caddyhttp"
    14  
    15  	"go.uber.org/zap"
    16  
    17  	"github.com/imgk/caddy-trojan/app"
    18  	"github.com/imgk/caddy-trojan/trojan"
    19  	"github.com/imgk/caddy-trojan/utils"
    20  	"github.com/imgk/caddy-trojan/websocket"
    21  )
    22  
    23  func init() {
    24  	caddy.RegisterModule(Handler{})
    25  	httpcaddyfile.RegisterHandlerDirective("trojan", func(h httpcaddyfile.Helper) (caddyhttp.MiddlewareHandler, error) {
    26  		m := &Handler{}
    27  		err := m.UnmarshalCaddyfile(h.Dispenser)
    28  		return m, err
    29  	})
    30  }
    31  
    32  // Handler implements an HTTP handler that ...
    33  type Handler struct {
    34  	WebSocket bool `json:"websocket,omitempty"`
    35  	Connect   bool `json:"connect_method,omitempty"`
    36  	Verbose   bool `json:"verbose,omitempty"`
    37  
    38  	// Upstream is ...
    39  	Upstream app.Upstream `json:"-,omitempty"`
    40  	// Proxy is ...
    41  	Proxy app.Proxy `json:"-,omitempty"`
    42  	// Logger is ...
    43  	Logger *zap.Logger `json:"-,omitempty"`
    44  	// Upgrader is ...
    45  	Upgrader websocket.Upgrader `json:"-,omitempty"`
    46  }
    47  
    48  // CaddyModule returns the Caddy module information.
    49  func (Handler) CaddyModule() caddy.ModuleInfo {
    50  	return caddy.ModuleInfo{
    51  		ID:  "http.handlers.trojan",
    52  		New: func() caddy.Module { return new(Handler) },
    53  	}
    54  }
    55  
    56  // Provision implements caddy.Provisioner.
    57  func (m *Handler) Provision(ctx caddy.Context) error {
    58  	m.Logger = ctx.Logger(m)
    59  	if !ctx.AppIsConfigured(app.CaddyAppID) {
    60  		return errors.New("handler: trojan is not configured")
    61  	}
    62  	mod, err := ctx.App(app.CaddyAppID)
    63  	if err != nil {
    64  		return err
    65  	}
    66  	app := mod.(*app.App)
    67  	m.Upstream = app.Upstream()
    68  	m.Proxy = app.Proxy()
    69  	return nil
    70  }
    71  
    72  // ServeHTTP implements caddyhttp.MiddlewareHandler.
    73  func (m *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request, next caddyhttp.Handler) error {
    74  	// trojan over http2/http3
    75  	// use CONNECT method, put trojan header as Proxy-Authorization
    76  	if m.Connect && r.Method == http.MethodConnect {
    77  		// handle trojan over http2/http3
    78  		if r.ProtoMajor == 1 {
    79  			return next.ServeHTTP(w, r)
    80  		}
    81  		auth := strings.TrimPrefix(r.Header.Get("Proxy-Authorization"), "Basic ")
    82  		if len(auth) != trojan.HeaderLen {
    83  			return next.ServeHTTP(w, r)
    84  		}
    85  		if ok := m.Upstream.Validate(auth); !ok {
    86  			return next.ServeHTTP(w, r)
    87  		}
    88  		if m.Verbose {
    89  			m.Logger.Info(fmt.Sprintf("handle trojan http%d from %v", r.ProtoMajor, r.RemoteAddr))
    90  		}
    91  
    92  		nr, nw, err := m.Proxy.Handle(r.Body, NewFlushWriter(w))
    93  		if err != nil {
    94  			m.Logger.Error(fmt.Sprintf("handle http%d error: %v", r.ProtoMajor, err))
    95  		}
    96  		m.Upstream.Consume(auth, nr, nw)
    97  		return nil
    98  	}
    99  
   100  	// handle websocket
   101  	if m.WebSocket && websocket.IsWebSocketUpgrade(r) {
   102  		conn, err := m.Upgrader.Upgrade(w, r, nil)
   103  		if err != nil {
   104  			return err
   105  		}
   106  
   107  		c := websocket.NewConn(conn)
   108  		defer c.Close()
   109  
   110  		b := [trojan.HeaderLen + 2]byte{}
   111  		if _, err := io.ReadFull(c, b[:]); err != nil {
   112  			m.Logger.Error(fmt.Sprintf("read trojan header error: %v", err))
   113  			return nil
   114  		}
   115  		if ok := m.Upstream.Validate(utils.ByteSliceToString(b[:trojan.HeaderLen])); !ok {
   116  			return nil
   117  		}
   118  		if m.Verbose {
   119  			m.Logger.Info(fmt.Sprintf("handle trojan websocket.Conn from %v", r.RemoteAddr))
   120  		}
   121  
   122  		nr, nw, err := m.Proxy.Handle(io.Reader(c), io.Writer(c))
   123  		if err != nil {
   124  			m.Logger.Error(fmt.Sprintf("handle websocket error: %v", err))
   125  		}
   126  		m.Upstream.Consume(utils.ByteSliceToString(b[:trojan.HeaderLen]), nr, nw)
   127  		return nil
   128  	}
   129  	return next.ServeHTTP(w, r)
   130  }
   131  
   132  // UnmarshalCaddyfile unmarshals Caddyfile tokens into h.
   133  func (h *Handler) UnmarshalCaddyfile(d *caddyfile.Dispenser) error {
   134  	if !d.Next() {
   135  		return d.ArgErr()
   136  	}
   137  	args := d.RemainingArgs()
   138  	if len(args) > 0 {
   139  		return d.ArgErr()
   140  	}
   141  	for nesting := d.Nesting(); d.NextBlock(nesting); {
   142  		subdirective := d.Val()
   143  		switch subdirective {
   144  		case "websocket":
   145  			if h.WebSocket {
   146  				return d.Err("only one websocket is not allowed")
   147  			}
   148  			h.WebSocket = true
   149  		case "connect_method":
   150  			if h.Connect {
   151  				return d.Err("only one connect_method is not allowed")
   152  			}
   153  			h.Connect = true
   154  		case "verbose":
   155  			if h.Verbose {
   156  				return d.Err("only one verbose is not allowed")
   157  			}
   158  			h.Verbose = true
   159  		}
   160  	}
   161  	return nil
   162  }
   163  
   164  // Interface guards
   165  var (
   166  	_ caddy.Provisioner           = (*Handler)(nil)
   167  	_ caddyhttp.MiddlewareHandler = (*Handler)(nil)
   168  	_ caddyfile.Unmarshaler       = (*Handler)(nil)
   169  )
   170  
   171  // FlushWriter is ...
   172  type FlushWriter struct {
   173  	Writer  io.Writer
   174  	Flusher http.Flusher
   175  }
   176  
   177  // NewFlushWriter is ...
   178  func NewFlushWriter(w http.ResponseWriter) *FlushWriter {
   179  	return &FlushWriter{
   180  		Writer:  w,
   181  		Flusher: w.(http.Flusher),
   182  	}
   183  }
   184  
   185  // Write is ...
   186  func (c *FlushWriter) Write(b []byte) (int, error) {
   187  	n, err := c.Writer.Write(b)
   188  	c.Flusher.Flush()
   189  	return n, err
   190  }