Merge branch 'feature/cache-enctable' into develop

cmd/shadowsocks-local/local.go

@@ -196,20 +196,6 @@ func enoughOptions(config *ss.Config) bool {
 		config.LocalPort != 0 && config.Password != ""
 }
 
-func isFileExists(path string) (bool, error) {
-	stat, err := os.Stat(path)
-	if err == nil {
-		if stat.Mode()&os.ModeType == 0 {
-			return true, nil
-		}
-		return false, errors.New(path + " exists but is not regular file")
-	}
-	if os.IsNotExist(err) {
-		return false, nil
-	}
-	return false, err
-}
-
 func main() {
 	var configFile string
 	var cmdConfig ss.Config
@@ -223,7 +209,7 @@ func main() {
 
 	flag.Parse()
 
-	exists, err := isFileExists(configFile)
+	exists, err := ss.IsFileExists(configFile)
 	// If no config file in current directory, try search it in the binary directory
 	// Note there's no portable way to detect the binary directory.
 	binDir := path.Dir(os.Args[0])

cmd/shadowsocks-server/server.go

@@ -1,8 +1,10 @@
 package main
 
 import (
+	"bufio"
 	"bytes"
 	"encoding/binary"
+	"encoding/gob"
 	"errors"
 	"flag"
 	ss "github.com/shadowsocks/shadowsocks-go/shadowsocks"
@@ -125,22 +127,88 @@ func handleConnection(conn *ss.Conn) {
 	return
 }
 
-// Add a encrypt table cache to save memory and startup time in case of many
-// same password.
-// If startup time becomes an issue, save the encrypt table on disk.
-var tableCache = map[string]*ss.EncryptTable{}
-var tableGetCnt int32
+const tableCacheFile = "table.cache"
+
+var table struct {
+	cache  map[string]*ss.EncryptTable
+	getCnt int32
+	hitCnt int32
+}
+
+func initTableCache(config *ss.Config) {
+	var exists bool
+	var err error
+	if !config.CacheEncTable {
+		goto emptyCache
+	}
+	exists, err = ss.IsFileExists(tableCacheFile)
+	if exists {
+		// load table cache from file
+		f, err := os.Open(tableCacheFile)
+		if err != nil {
+			log.Println("error opening table cache:", err)
+			goto emptyCache
+		}
+		defer f.Close()
+
+		dec := gob.NewDecoder(bufio.NewReader(f))
+		if err = dec.Decode(&table.cache); err != nil {
+			log.Println("error loading table cache:", err)
+			goto emptyCache
+		}
+		debug.Println("table cache loaded from disk")
+		return
+	}
+	if err != nil {
+		log.Println("table cache:", err)
+	}
+
+emptyCache:
+	debug.Println("creating empty table cache")
+	table.cache = map[string]*ss.EncryptTable{}
+}
+
+func storeTableCache(config *ss.Config) {
+	if !config.CacheEncTable || table.getCnt == table.hitCnt {
+		return
+	}
+
+	const tmpPath = "tmp.cache"
+	f, err := os.Create(tmpPath)
+	if err != nil {
+		log.Println("can't create tmp table cache")
+		return
+	}
+	defer f.Close()
+
+	w := bufio.NewWriter(f)
+	enc := gob.NewEncoder(w)
+	if err = enc.Encode(table.cache); err != nil {
+		log.Println("error saving tmp table cache:", err)
+		return
+	}
+	if err = w.Flush(); err != nil {
+		log.Println("error flushing table cache:", err)
+		return
+	}
+	if err = os.Rename(tmpPath, tableCacheFile); err != nil {
+		log.Printf("error renaming %s to %s: %v\n", tmpPath, tableCacheFile, err)
+		return
+	}
+	debug.Println("table cache saved")
+}
 
 func getTable(password string) (tbl *ss.EncryptTable) {
-	if tableCache != nil {
+	if table.cache != nil {
 		var ok bool
-		tbl, ok = tableCache[password]
+		tbl, ok = table.cache[password]
 		if ok {
+			atomic.AddInt32(&table.hitCnt, 1)
 			debug.Println("table cache hit for password:", password)
 			return
 		}
 		tbl = ss.GetTable(password)
-		tableCache[password] = tbl
+		table.cache[password] = tbl
 	} else {
 		tbl = ss.GetTable(password)
 	}
@@ -248,7 +316,7 @@ func run(port, password string) {
 	}
 	passwdManager.add(port, password, ln)
 	encTbl := getTable(password)
-	atomic.AddInt32(&tableGetCnt, 1)
+	atomic.AddInt32(&table.getCnt, 1)
 	log.Printf("server listening port %v ...\n", port)
 	for {
 		conn, err := ln.Accept()
@@ -294,6 +362,7 @@ func main() {
 	flag.BoolVar((*bool)(&debug), "d", false, "print debug message")
 
 	flag.Parse()
+	ss.SetDebug(debug)
 
 	var err error
 	config, err = ss.ParseConfig(configFile)
@@ -310,19 +379,22 @@ func main() {
 	} else {
 		ss.UpdateConfig(config, &cmdConfig)
 	}
-	ss.SetDebug(debug)
 
 	if err = unifyPortPassword(config); err != nil {
 		os.Exit(1)
 	}
+
+	initTableCache(config)
 	for port, password := range config.PortPassword {
 		go run(port, password)
 	}
 	// Wait all ports have get it's encryption table
-	for int(tableGetCnt) != len(config.PortPassword) {
+	for int(table.getCnt) != len(config.PortPassword) {
 		time.Sleep(1 * time.Second)
 	}
+	storeTableCache(config)
 	log.Println("all ports ready")
-	tableCache = nil // release memory
+
+	table.cache = nil // release memory
 	waitSignal()
 }

config.json

@@ -7,5 +7,6 @@
         "8388": "barfoo!",
         "8387": "foobar!"
     },
-    "timeout":60
+    "timeout":60,
+    "cache_enctable":false
 }

shadowsocks/config.go

@@ -16,12 +16,13 @@ import (
 )
 
 type Config struct {
-	Server       string            `json:"server"`
-	ServerPort   int               `json:"server_port"`
-	LocalPort    int               `json:"local_port"`
-	Password     string            `json:"password"`
-	PortPassword map[string]string `json:"port_password"`
-	Timeout      int               `json:"timeout"`
+	Server        string            `json:"server"`
+	ServerPort    int               `json:"server_port"`
+	LocalPort     int               `json:"local_port"`
+	Password      string            `json:"password"`
+	PortPassword  map[string]string `json:"port_password"`
+	Timeout       int               `json:"timeout"`
+	CacheEncTable bool              `json:"cache_enctable"`
 }
 
 var readTimeout time.Duration

shadowsocks/config_test.go

@@ -5,21 +5,27 @@ import (
 )
 
 func TestParseConfig1Password(t *testing.T) {
-	config := ParseConfig("testdata/config-one-passwd.json")
+	config, err := ParseConfig("testdata/config-one-passwd.json")
+	if err != nil {
+		t.Error("error parsing single password config:", err)
+	}
 
 	if config.Password != "barfoo!" {
 		t.Error("wrong password from config")
 	}
-	if config.Debug != true {
-		t.Error("debug option wrong")
-	}
 	if config.Timeout != 60 {
 		t.Error("tiemout wrong")
 	}
+	if !config.CacheEncTable {
+		t.Error("cache_enctable should be true")
+	}
 }
 
 func TestParseConfigMultiPassword(t *testing.T) {
-	config := ParseConfig("testdata/config-multi-passwd.json")
+	config, err := ParseConfig("testdata/config-multi-passwd.json")
+	if err != nil {
+		t.Error("error parsing multi password config:", err)
+	}
 
 	if config.Password != "barfoo!" {
 		t.Error("wrong password from config")

shadowsocks/conn.go

@@ -70,13 +70,13 @@ func (c Conn) Read(b []byte) (n int, err error) {
 	buf := make([]byte, len(b), len(b))
 	n, err = c.Conn.Read(buf)
 	if n > 0 {
-		encrypt2(c.decTbl, buf[0:n], b[0:n])
+		encrypt2(c.DecTbl, buf[0:n], b[0:n])
 	}
 	return
 }
 
 func (c Conn) Write(b []byte) (n int, err error) {
-	buf := encrypt(c.encTbl, b)
+	buf := encrypt(c.EncTbl, b)
 	n, err = c.Conn.Write(buf)
 	return
 }

shadowsocks/encrypt.go

@@ -8,8 +8,8 @@ import (
 )
 
 type EncryptTable struct {
-	encTbl []byte
-	decTbl []byte
+	EncTbl []byte
+	DecTbl []byte
 }
 
 func GetTable(key string) (tbl *EncryptTable) {
@@ -38,10 +38,10 @@ func GetTable(key string) (tbl *EncryptTable) {
 		})
 	}
 	for i = 0; i < tbl_size; i++ {
-		tbl.encTbl[i] = byte(table[i])
+		tbl.EncTbl[i] = byte(table[i])
 	}
 	for i = 0; i < tbl_size; i++ {
-		tbl.decTbl[tbl.encTbl[i]] = byte(i)
+		tbl.DecTbl[tbl.EncTbl[i]] = byte(i)
 	}
 	return
 }

shadowsocks/encrypt_test.go

@@ -8,10 +8,10 @@ const tbl_size = 256
 
 func checkTable(t *testing.T, tbl *EncryptTable, encTarget, decTarget []byte, msg string) {
 	for i := 0; i < tbl_size; i++ {
-		if encTarget[i] != tbl.encTbl[i] {
+		if encTarget[i] != tbl.EncTbl[i] {
 			t.Fatalf("%s: encrypt table error at index %d\n", msg, i)
 		}
-		if decTarget[i] != tbl.decTbl[i] {
+		if decTarget[i] != tbl.DecTbl[i] {
 			t.Fatalf("%s: decrypt table error at index %d\n", msg, i)
 		}
 	}

shadowsocks/testdata/config-one-passwd.json

@@ -4,5 +4,5 @@
     "local_port":1081,
     "password":"barfoo!",
     "timeout":60,
-    "debug": true
+    "cache_enctable": true
 }

shadowsocks/util.go

@@ -0,0 +1,20 @@
+package shadowsocks
+
+import (
+	"errors"
+	"os"
+)
+
+func IsFileExists(path string) (bool, error) {
+	stat, err := os.Stat(path)
+	if err == nil {
+		if stat.Mode()&os.ModeType == 0 {
+			return true, nil
+		}
+		return false, errors.New(path + " exists but is not regular file")
+	}
+	if os.IsNotExist(err) {
+		return false, nil
+	}
+	return false, err
+}