Reimplement transocks based on cybozu-go/cmd .

.travis.yml

@@ -1,8 +1,14 @@
 sudo: false
 language: go
 go:
-  - 1.6
+  - 1.7
   - tip
 
+before_install:
+  - go get github.com/golang/lint/golint
+
 script:
+  - go install ./...
   - go test -v ./...
+  - go vet -x ./...
+  - $HOME/gopath/bin/golint -set_exit_status -min_confidence 0.81 ./...

CHANGELOG.md

@@ -0,0 +1,16 @@
+# Change Log
+
+All notable changes to this project will be documented in this file.
+
+## [Unreleased]
+### Added
+- transocks now adopts [github.com/cybozu-go/cmd][cmd] framework.  
+  As a result, it implements [the common spec][spec] including graceful restart.
+
+### Changed
+- The default configuration file path is now `/etc/transocks.toml`.
+- Configuration items for logging is changed.
+
+[cmd]: https://github.com/cybozu-go/cmd
+[spec]: https://github.com/cybozu-go/cmd/blob/master/README.md#specifications
+[Unreleased]: https://github.com/cybozu-go/transocks/compare/v0.1...HEAD

README.md

@@ -1,5 +1,8 @@
-[![GoDoc](https://godoc.org/github.com/cybozu-go/transocks?status.png)][godoc]
+[![GitHub release](https://img.shields.io/github/release/cybozu-go/transocks.svg?maxAge=60)][releases]
-[![Build Status](https://travis-ci.org/cybozu-go/transocks.png)](https://travis-ci.org/cybozu-go/transocks)
+[![GoDoc](https://godoc.org/github.com/cybozu-go/transocks?status.svg)][godoc]
+[![Build Status](https://travis-ci.org/cybozu-go/transocks.svg?branch=master)](https://travis-ci.org/cybozu-go/transocks)
+[![Go Report Card](https://goreportcard.com/badge/github.com/cybozu-go/transocks)](https://goreportcard.com/report/github.com/cybozu-go/transocks)
+[![License](https://img.shields.io/github/license/cybozu-go/transocks.svg?maxAge=2592000)](LICENSE)
 
 transocks - a transparent SOCKS5/HTTP proxy
 ===========================================
@@ -21,35 +24,42 @@ Features
 * SOCKS5 and HTTP proxy (CONNECT)
 
     We recommend using SOCKS5 server if available.
-    Looking for a good SOCKS5 server?  Take a look at our [usocksd][]!
+    Take a look at our SOCKS server [usocksd][] if you are looking for.
 
     HTTP proxies often prohibits CONNECT method to make connections
     to ports other than 443.  Make sure your HTTP proxy allows CONNECT
     to the ports you want.
 
+* Graceful stop & restart
+
+    * On SIGINT/SIGTERM, transocks stops gracefully.
+    * On SIGHUP, transocks restarts gracefully.
+
 * Library and executable
 
     transocks comes with a handy executable.
     You may use the library to create your own.
 
+Install
+-------
+
+Use Go 1.7 or better.
+
+```
+go get -u github.com/cybozu-go/transocks/...
+```
+
 Usage
 -----
 
 `transocks [-h] [-f CONFIG]`
 
-The default configuration file path is `/usr/local/etc/transocks.toml`.
+The default configuration file path is `/etc/transocks.toml`.
 
-`transocks` does not have *daemon* mode.  Use systemd or upstart to
+In addition, transocks implements [the common spec](https://github.com/cybozu-go/cmd#specifications) from [`cybozu-go/cmd`](https://github.com/cybozu-go/cmd).
-run it on your background.
 
-Install
+transocks does not have *daemon* mode.  Use systemd to run it
--------
+on your background.
-
-Use Go 1.5 or better.
-
-```
-go get github.com/cybozu-go/transocks/cmd/transocks
-```
 
 Configuration file format
 -------------------------
@@ -66,8 +76,10 @@ listen = "localhost:1081"
 proxy_url = "socks5://10.20.30.40:1080"  # for SOCKS5 server
 #proxy_url = "http://10.20.30.40:3128"   # for HTTP proxy server
 
-log_level = "info"
+[log]
-log_file = "/var/log/transocks.log"
+filename = "/path/to/file"   # default to stderr
+level = "info"               # critical", error, warning, info, debug
+format = "json"              # plain, logfmt, json
 ```
 
 Redirecting connections by iptables
@@ -110,13 +122,7 @@ License
 
 [MIT](https://opensource.org/licenses/MIT)
 
-Author
-------
-
-[@ymmt2005][]
-
 [godoc]: https://godoc.org/github.com/cybozu-go/transocks
 [Squid]: http://www.squid-cache.org/
 [usocksd]: https://github.com/cybozu-go/usocksd
 [TOML]: https://github.com/toml-lang/toml
-[@ymmt2005]: https://github.com/ymmt2005

cmd/transocks/main.go

@@ -1,79 +1,90 @@
-// transocks server.
 package main
 
 import (
 	"flag"
 	"fmt"
+	"net"
 	"net/url"
-	"os"
 
 	"github.com/BurntSushi/toml"
+	"github.com/cybozu-go/cmd"
 	"github.com/cybozu-go/log"
 	"github.com/cybozu-go/transocks"
 )
 
 type tomlConfig struct {
-	Listen   string
+	Listen   string        `toml:"listen"`
-	ProxyURL string `toml:"proxy_url"`
+	ProxyURL string        `toml:"proxy_url"`
-	LogLevel string `toml:"log_level"`
+	Log      cmd.LogConfig `toml:"log"`
-	LogFile  string `toml:"log_file"`
 }
 
 var (
-	configFile = flag.String("f", "/usr/local/etc/transocks.toml",
+	configFile = flag.String("f", "/etc/transocks.toml",
 		"TOML configuration file path")
 )
 
-func loadConfig() (*transocks.Config, string, error) {
+func loadConfig() (*transocks.Config, error) {
 	tc := new(tomlConfig)
 	md, err := toml.DecodeFile(*configFile, tc)
 	if err != nil {
-		return nil, "", err
+		return nil, err
 	}
 	if len(md.Undecoded()) > 0 {
-		return nil, "", fmt.Errorf("undecoded key in TOML: %v", md.Undecoded())
+		return nil, fmt.Errorf("undecoded key in TOML: %v", md.Undecoded())
 	}
 
 	c := transocks.NewConfig()
-	c.Listen = tc.Listen
+	c.Addr = tc.Listen
 
 	u, err := url.Parse(tc.ProxyURL)
 	if err != nil {
-		return nil, "", err
+		return nil, err
 	}
 	c.ProxyURL = u
 
-	if err = log.DefaultLogger().SetThresholdByName(tc.LogLevel); err != nil {
+	err = tc.Log.Apply()
-		return nil, "", err
+	if err != nil {
+		return nil, err
 	}
 
-	return c, tc.LogFile, nil
+	return c, nil
+}
+
+func serve(lns []net.Listener, c *transocks.Config) {
+	s, err := transocks.NewServer(c)
+	if err != nil {
+		log.ErrorExit(err)
+	}
+
+	for _, ln := range lns {
+		s.Serve(ln)
+	}
+	err = cmd.Wait()
+	if err != nil && !cmd.IsSignaled(err) {
+		log.ErrorExit(err)
+	}
 }
 
 func main() {
 	flag.Parse()
 
-	c, logfile, err := loadConfig()
+	c, err := loadConfig()
 	if err != nil {
 		log.ErrorExit(err)
 	}
 
-	if len(logfile) > 0 {
+	g := &cmd.Graceful{
-		f, err := os.OpenFile(logfile, os.O_WRONLY|os.O_APPEND, 0644)
+		Listen: func() ([]net.Listener, error) {
-		if err != nil {
+			return transocks.Listeners(c)
-			log.ErrorExit(err)
+		},
-		}
+		Serve: func(lns []net.Listener) {
-		defer f.Close()
+			serve(lns, c)
-		log.DefaultLogger().SetOutput(f)
+		},
 	}
+	g.Run()
 
-	srv, err := transocks.NewServer(c)
+	err = cmd.Wait()
-	if err != nil {
+	if err != nil && !cmd.IsSignaled(err) {
 		log.ErrorExit(err)
 	}
-	log.Info("server starts", nil)
-
-	srv.Serve()
-
-	log.Info("server ends", nil)
 }

config.go

@@ -5,18 +5,32 @@ import (
 	"fmt"
 	"net"
 	"net/url"
+	"time"
+
+	"github.com/cybozu-go/cmd"
+	"github.com/cybozu-go/log"
 )
 
 const (
-	// NAT mode
+	defaultShutdownTimeout = 1 * time.Minute
-	ModeNAT = "nat"
+)
+
+// Mode is the type of transocks mode.
+type Mode string
+
+func (m Mode) String() string {
+	return string(m)
+}
+
+const (
+	// ModeNAT is mode constant for NAT.
+	ModeNAT = Mode("nat")
 )
 
 // Config keeps configurations for Server.
 type Config struct {
-	// Listen is the listening address.
+	// Addr is the listening address.
-	// e.g. "localhost:1081"
+	Addr string
-	Listen string
 
 	// ProxyURL is the URL for upstream proxy.
 	//
@@ -27,27 +41,39 @@ type Config struct {
 	ProxyURL *url.URL
 
 	// Mode determines how clients are routed to transocks.
-	// Default is "nat".  No other options are available at this point.
+	// Default is ModeNAT.  No other options are available at this point.
-	Mode string
+	Mode Mode
+
+	// ShutdownTimeout is the maximum duration the server waits for
+	// all connections to be closed before shutdown.
+	//
+	// Zero duration disables timeout.  Default is 1 minute.
+	ShutdownTimeout time.Duration
 
 	// Dialer is the base dialer to connect to the proxy server.
 	// The server uses the default dialer if this is nil.
 	Dialer *net.Dialer
+
+	// Logger can be used to provide a custom logger.
+	// If nil, the default logger is used.
+	Logger *log.Logger
+
+	// Env can be used to specify a cmd.Environment on which the server runs.
+	// If nil, the server will run on the global environment.
+	Env *cmd.Environment
 }
 
 // NewConfig creates and initializes a new Config.
 func NewConfig() *Config {
 	c := new(Config)
 	c.Mode = ModeNAT
+	c.ShutdownTimeout = defaultShutdownTimeout
 	return c
 }
 
 // validate validates the configuration.
 // It returns non-nil error if the configuration is not valid.
 func (c *Config) validate() error {
-	if len(c.Listen) == 0 {
-		return errors.New("Listen is empty")
-	}
 	if c.ProxyURL == nil {
 		return errors.New("ProxyURL is nil")
 	}

original_dst_linux.go

@@ -10,7 +10,10 @@ import (
 )
 
 const (
-	SO_ORIGINAL_DST      = 80
+	// SO_ORIGINAL_DST is a Linux getsockopt optname.
+	SO_ORIGINAL_DST = 80
+
+	// IP6T_SO_ORIGINAL_DST a Linux getsockopt optname.
 	IP6T_SO_ORIGINAL_DST = 80
 )
 

original_dst_linux_test.go

@@ -22,10 +22,10 @@ func TestGetOriginalDST(t *testing.T) {
 	}
 	defer c.Close()
 
-	orig_addr, err := GetOriginalDST(c.(*net.TCPConn))
+	origAddr, err := GetOriginalDST(c.(*net.TCPConn))
 	if err != nil {
 		t.Fatal(err)
 	}
 
-	t.Log(orig_addr.String())
+	t.Log(origAddr.String())
 }

server.go

@@ -1,26 +1,39 @@
 package transocks
 
 import (
+	"context"
 	"io"
 	"net"
+	"sync"
 	"time"
 
+	"github.com/cybozu-go/cmd"
 	"github.com/cybozu-go/log"
+	"github.com/cybozu-go/netutil"
 	"golang.org/x/net/proxy"
 )
 
-var (
+const (
-	defaultDialer = &net.Dialer{
+	keepAliveTimeout = 3 * time.Minute
-		Timeout:   10 * time.Second,
+	copyBufferSize   = 64 << 10
-		KeepAlive: 60 * time.Second,
-	}
 )
 
+// Listeners returns a list of net.Listener.
+func Listeners(c *Config) ([]net.Listener, error) {
+	ln, err := net.Listen("tcp", c.Addr)
+	if err != nil {
+		return nil, err
+	}
+	return []net.Listener{ln}, nil
+}
+
 // Server provides transparent proxy server functions.
 type Server struct {
-	config   *Config
+	cmd.Server
-	dialer   proxy.Dialer
+	mode   Mode
-	listener net.Listener
+	logger *log.Logger
+	dialer proxy.Dialer
+	pool   sync.Pool
 }
 
 // NewServer creates Server.
@@ -30,99 +43,110 @@ func NewServer(c *Config) (*Server, error) {
 		return nil, err
 	}
 
-	dialer := defaultDialer
+	dialer := c.Dialer
-	if c.Dialer != nil {
+	if dialer == nil {
-		dialer = c.Dialer
+		dialer = &net.Dialer{
+			KeepAlive: keepAliveTimeout,
+			DualStack: true,
+		}
 	}
-	proxy_dialer, err := proxy.FromURL(c.ProxyURL, dialer)
+	pdialer, err := proxy.FromURL(c.ProxyURL, dialer)
 	if err != nil {
 		return nil, err
 	}
-
+	logger := c.Logger
-	l, err := net.Listen("tcp", c.Listen)
+	if logger == nil {
-	if err != nil {
+		logger = log.DefaultLogger()
-		return nil, err
 	}
 
-	return &Server{c, proxy_dialer, l}, nil
+	s := &Server{
+		Server: cmd.Server{
+			ShutdownTimeout: c.ShutdownTimeout,
+			Env:             c.Env,
+		},
+		mode:   c.Mode,
+		logger: logger,
+		dialer: pdialer,
+		pool: sync.Pool{
+			New: func() interface{} {
+				return make([]byte, copyBufferSize)
+			},
+		},
+	}
+	s.Server.Handler = s.handleConnection
+	return s, nil
 }
 
-// Serve accepts and handles new connections forever.
+func (s *Server) handleConnection(ctx context.Context, conn net.Conn) {
-func (s *Server) Serve() error {
+	tc, ok := conn.(*net.TCPConn)
-	for {
+	if !ok {
-		conn, err := s.listener.Accept()
+		s.logger.Error("non-TCP connection", map[string]interface{}{
-		if err != nil {
+			"conn": conn,
-			log.Critical(err.Error(), nil)
-			return err
-		}
-		tcp_conn, ok := conn.(*net.TCPConn)
-		if !ok {
-			conn.Close()
-			panic("not a TCPConn!")
-		}
-		go s.handleConnection(tcp_conn)
-	}
-}
-
-func (s *Server) handleConnection(c *net.TCPConn) {
-	defer c.Close()
-
-	var addr string
-
-	switch s.config.Mode {
-	case ModeNAT:
-		orig_addr, err := GetOriginalDST(c)
-		if err != nil {
-			log.Error(err.Error(), nil)
-			return
-		}
-		addr = orig_addr.String()
-	default:
-		addr = c.LocalAddr().String()
-	}
-
-	if log.Enabled(log.LvDebug) {
-		log.Debug("making proxy connection", map[string]interface{}{
-			"_dst": addr,
-		})
-	}
-
-	pconn, err := s.dialer.Dial("tcp", addr)
-	if err != nil {
-		log.Error(err.Error(), map[string]interface{}{
-			"_dst": addr,
 		})
 		return
 	}
-	defer pconn.Close()
 
-	ch := make(chan error, 2)
+	fields := cmd.FieldsFromContext(ctx)
-	go copyData(c, pconn, ch)
+	fields[log.FnType] = "access"
-	go copyData(pconn, c, ch)
+	fields["client_addr"] = conn.RemoteAddr().String()
-	for i := 0; i < 2; i++ {
+
-		e := <-ch
+	var addr string
-		if e != nil {
+	switch s.mode {
-			log.Error(e.Error(), map[string]interface{}{
+	case ModeNAT:
-				"_dst": addr,
+		origAddr, err := GetOriginalDST(tc)
-			})
+		if err != nil {
-			break
+			fields[log.FnError] = err.Error()
+			s.logger.Error("GetOriginalDST failed", fields)
+			return
 		}
+		addr = origAddr.String()
+	default:
+		addr = tc.LocalAddr().String()
 	}
+	fields["dest_addr"] = addr
 
-	if log.Enabled(log.LvDebug) {
+	destConn, err := s.dialer.Dial("tcp", addr)
-		log.Debug("closing proxy connection", map[string]interface{}{
+	if err != nil {
-			"_dst": addr,
+		fields[log.FnError] = err.Error()
-		})
+		s.logger.Error("failed to connect to proxy server", fields)
+		return
 	}
-}
+	defer destConn.Close()
 
-func copyData(dst io.Writer, src io.Reader, ch chan<- error) {
+	s.logger.Info("proxy starts", fields)
-	_, err := io.Copy(dst, src)
+
-	if tdst, ok := dst.(*net.TCPConn); ok {
+	// do proxy
-		tdst.CloseWrite()
+	st := time.Now()
+	env := cmd.NewEnvironment(ctx)
+	env.Go(func(ctx context.Context) error {
+		buf := s.pool.Get().([]byte)
+		_, err := io.CopyBuffer(destConn, tc, buf)
+		s.pool.Put(buf)
+		if hc, ok := destConn.(netutil.HalfCloser); ok {
+			hc.CloseWrite()
+		}
+		tc.CloseRead()
+		return err
+	})
+	env.Go(func(ctx context.Context) error {
+		buf := s.pool.Get().([]byte)
+		_, err := io.CopyBuffer(tc, destConn, buf)
+		s.pool.Put(buf)
+		tc.CloseWrite()
+		if hc, ok := destConn.(netutil.HalfCloser); ok {
+			hc.CloseRead()
+		}
+		return err
+	})
+	env.Stop()
+	err = env.Wait()
+
+	fields = cmd.FieldsFromContext(ctx)
+	fields["elapsed"] = time.Since(st).Seconds()
+	if err != nil {
+		fields[log.FnError] = err.Error()
+		s.logger.Error("proxy ends with an error", fields)
+		return
 	}
-	if tsrc, ok := src.(*net.TCPConn); ok {
+	s.logger.Info("proxy ends", fields)
-		tsrc.CloseRead()
-	}
-	ch <- err
 }