Bug fixes.

README.md

@@ -22,3 +22,7 @@ When you want to make KumaSocks more smaller, you can consider to use upx. Here
 ```
 upx --best KumaSocks
 ```
+
+## FAQS
+
+If you have problems with splice in Linux, try to turn on the `io-copy-hack` option in config file. It will reduce the performance as it is using a custom version of io.Copy without using splice system call.

kumasocks-sample.toml

@@ -1,3 +1,3 @@
 listen-addr = "0.0.0.0:1234"
 proxy-addr = "socks5://192.168.1.210:1080"
-
+io-copy-hack = false

main.go

@@ -3,15 +3,13 @@ package main
 
 import (
 	"flag"
-	"fmt"
 	"io"
 	"io/ioutil"
 	"log"
 	"net"
 	"net/url"
 	"os"
 	"os/signal"
-	"sync"
 	"syscall"
 	"time"
 
@@ -23,60 +21,102 @@ var (
 	configPath  string
 	conf        Config
 	proxyDialer proxy.Dialer
+	running     bool
 )
 
 type Config struct {
 	ListenAddr string `toml:"listen-addr"`
 	ProxyAddr  string `toml:"proxy-addr"`
+	IOCopyHack bool   `toml:"io-copy-hack"`
 }
 
 func readConf(path string) string {
 	buf, err := ioutil.ReadFile(path)
 	if err != nil {
-		log.Fatal(err)
-		os.Exit(0)
+		log.Fatalf("[ERROR] Read config error: %s.\n", err.Error())
 		return ""
 	}
 
 	return string(buf)
 }
 
+func customCopy(dst io.Writer, src io.Reader, buf []byte) (err error) {
+	for {
+		nr, er := src.Read(buf)
+		if nr > 0 {
+			_, ew := dst.Write(buf[0:nr])
+			if ew != nil {
+				err = ew
+				break
+			}
+		}
+		if er != nil {
+			if er != io.EOF {
+				err = er
+			}
+			break
+		}
+	}
+
+	return err
+}
+
 func handleConnection(conn *net.TCPConn) {
-	defer conn.Close()
+	conn.SetKeepAlive(true)
+	conn.SetKeepAlivePeriod(20 * time.Second)
+	conn.SetNoDelay(true)
 
 	dstAddr, err := GetOriginalDST(conn)
 	if err != nil {
-		log.Fatal(err)
+		log.Printf("[ERROR] GetOriginalDST error: %s\n", err.Error())
 		return
 	}
 
-	//log.Println("Incoming connection to " + dstAddr.String())
+	log.Printf("[INFO] Connect to %s.\n", dstAddr.String())
 
 	proxy, err := proxyDialer.Dial("tcp", dstAddr.String())
 	if err != nil {
-		log.Println(err)
+		log.Printf("[ERROR] Dial proxy error: %s.\n", err.Error())
 		return
 	}
-	defer proxy.Close()
 
-	wg := sync.WaitGroup{}
-	wg.Add(2)
+	proxy.(*net.TCPConn).SetKeepAlive(true)
+	proxy.(*net.TCPConn).SetKeepAlivePeriod(20 * time.Second)
+	proxy.(*net.TCPConn).SetNoDelay(true)
 
-	go func() {
-		io.Copy(proxy, conn)
-		conn.Close()
-		proxy.Close()
-		wg.Done()
-	}()
+	copyEnd := false
 
 	go func() {
-		io.Copy(conn, proxy)
+		buf := make([]byte, 32*1024)
+
+		var err error
+		if conf.IOCopyHack {
+			err = customCopy(proxy, conn, buf)
+		} else {
+			_, err = io.CopyBuffer(proxy, conn, buf)
+		}
+
+		if err != nil && !copyEnd {
+			log.Printf("[ERROR] Copy error: %s.\n", err.Error())
+		}
+		copyEnd = true
 		conn.Close()
 		proxy.Close()
-		wg.Done()
 	}()
 
-	wg.Wait()
+	buf := make([]byte, 32*1024)
+	if conf.IOCopyHack {
+		err = customCopy(conn, proxy, buf)
+	} else {
+		_, err = io.CopyBuffer(conn, proxy, buf)
+	}
+
+	if err != nil && !copyEnd {
+		log.Printf("[ERROR] Copy error: %s.\n", err.Error())
+	}
+	copyEnd = true
+	conn.Close()
+	proxy.Close()
 }
 
 func main() {
@@ -85,42 +125,42 @@ func main() {
 
 	configStr := readConf(configPath)
 	if _, err := toml.Decode(configStr, &conf); err != nil {
-		log.Fatal(err)
-		return
+		log.Fatalf("[ERROR] Decode config error: %s.\n", err.Error())
 	}
 
 	listener, err := net.Listen("tcp", conf.ListenAddr)
 	if err != nil {
-		log.Fatal(err)
-		return
+		log.Fatalf("[ERROR] Listen tcp error: %s.\n", err.Error())
 	}
 
 	defer listener.Close()
 
 	proxyURL, err := url.Parse(conf.ProxyAddr)
 	if err != nil {
-		log.Fatal(err)
-		return
+		log.Fatalf("[ERROR] Parse proxy address error: %s.\n", err.Error())
 	}
 
-	dialer := &net.Dialer{
-		KeepAlive: 30 * time.Second,
-		DualStack: true,
+	proxyDialer, err = proxy.FromURL(proxyURL, &net.Dialer{})
+	if err != nil {
+		log.Fatalf("[ERROR] Create proxy error: %s.\n", err.Error())
 	}
 
-	proxyDialer, err = proxy.FromURL(proxyURL, dialer)
-	if err != nil {
-		log.Fatal(err)
-		return
+	log.Println("[INFO] Starting Kumasocks...")
+
+	if conf.IOCopyHack {
+		log.Println("[INFO] Using io.Copy hack.")
 	}
 
-	fmt.Println("Starting Kumasocks...")
+	running = true
 
 	go func() {
-		for {
+		for running {
 			conn, err := listener.Accept()
 			if err != nil {
-				break
+				if running {
+					log.Printf("[ERROR] TCP accept error: %s.\n", err.Error())
+				}
+				continue
 			}
 			go handleConnection(conn.(*net.TCPConn))
 		}
@@ -130,7 +170,8 @@ func main() {
 	signal.Notify(c, syscall.SIGHUP, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT)
 	<-c
 
-	fmt.Println("Exiting Kumasocks...")
+	log.Println("[INFO] Exiting Kumasocks...")
 
+	running = false
 	listener.Close()
 }