From 7c315bf7dad5bba20d23b045591b9f259bf48912 Mon Sep 17 00:00:00 2001
From: notsure2 <notsure2@protonmail.com>
Date: Mon, 9 Dec 2019 18:12:47 +0200
Subject: [PATCH] Make keepalive optional on client -> server and server ->
 proxy connections. Use KeepAlive value in config (seconds).

---
 cmd/ck-client/ck-client.go |  4 ++--
 cmd/ck-server/ck-server.go |  3 ++-
 internal/client/state.go   |  7 +++++++
 internal/server/state.go   | 14 +++++++++++---
 4 files changed, 22 insertions(+), 6 deletions(-)

diff --git a/cmd/ck-client/ck-client.go b/cmd/ck-client/ck-client.go
index 7283303..8c654ee 100644
--- a/cmd/ck-client/ck-client.go
+++ b/cmd/ck-client/ck-client.go
@@ -23,7 +23,7 @@ import (
 var version string
 
 func makeSession(sta *client.State, isAdmin bool) *mux.Session {
-	log.Info("Attemtping to start a new session")
+	log.Info("Attempting to start a new session")
 	if !isAdmin {
 		// sessionID is usergenerated. There shouldn't be a security concern because the scope of
 		// sessionID is limited to its UID.
@@ -32,7 +32,7 @@ func makeSession(sta *client.State, isAdmin bool) *mux.Session {
 		atomic.StoreUint32(&sta.SessionID, binary.BigEndian.Uint32(quad))
 	}
 
-	d := net.Dialer{Control: protector}
+	d := net.Dialer{Control: protector, KeepAlive: sta.KeepAlive}
 	connsCh := make(chan net.Conn, sta.NumConn)
 	var _sessionKey atomic.Value
 	var wg sync.WaitGroup
diff --git a/cmd/ck-server/ck-server.go b/cmd/ck-server/ck-server.go
index 967a904..a58aede 100644
--- a/cmd/ck-server/ck-server.go
+++ b/cmd/ck-server/ck-server.go
@@ -174,7 +174,8 @@ func dispatchConnection(conn net.Conn, sta *server.State) {
 			}
 		}
 		proxyAddr := sta.ProxyBook[ci.ProxyMethod]
-		localConn, err := net.Dial(proxyAddr.Network(), proxyAddr.String())
+		d := net.Dialer{KeepAlive: sta.KeepAlive}
+		localConn, err := d.Dial(proxyAddr.Network(), proxyAddr.String())
 		if err != nil {
 			log.Errorf("Failed to connect to %v: %v", ci.ProxyMethod, err)
 			user.CloseSession(ci.SessionId, "Failed to connect to proxy server")
diff --git a/internal/client/state.go b/internal/client/state.go
index 14711d5..c96f141 100644
--- a/internal/client/state.go
+++ b/internal/client/state.go
@@ -24,6 +24,7 @@ type rawConfig struct {
 	Transport        string
 	NumConn          int
 	StreamTimeout    int
+	KeepAlive        int
 	RemoteHost       string
 	RemotePort       int
 }
@@ -50,6 +51,7 @@ type State struct {
 	ServerName       string
 	NumConn          int
 	Timeout          time.Duration
+	KeepAlive        time.Duration
 }
 
 // semi-colon separated value. This is for Android plugin options
@@ -138,6 +140,11 @@ func (sta *State) ParseConfig(conf string) (err error) {
 	} else {
 		sta.Timeout = time.Duration(preParse.StreamTimeout) * time.Second
 	}
+	if preParse.KeepAlive <= 0 {
+		sta.KeepAlive = -1
+	} else {
+		sta.KeepAlive = time.Duration(preParse.KeepAlive) * time.Second
+	}
 	sta.UID = preParse.UID
 
 	pub, ok := ecdh.Unmarshal(preParse.PublicKey)
diff --git a/internal/server/state.go b/internal/server/state.go
index 1ab2a50..2118039 100644
--- a/internal/server/state.go
+++ b/internal/server/state.go
@@ -24,6 +24,7 @@ type rawConfig struct {
 	AdminUID      []byte
 	DatabasePath  string
 	StreamTimeout int
+	KeepAlive     int
 	CncMode       bool
 }
 
@@ -32,9 +33,10 @@ type State struct {
 	BindAddr  []net.Addr
 	ProxyBook map[string]net.Addr
 
-	Now      func() time.Time
-	AdminUID []byte
-	Timeout  time.Duration
+	Now       func() time.Time
+	AdminUID  []byte
+	Timeout   time.Duration
+	KeepAlive time.Duration
 
 	BypassUID map[[16]byte]struct{}
 	staticPv  crypto.PrivateKey
@@ -173,6 +175,12 @@ func (sta *State) ParseConfig(conf string) (err error) {
 		sta.Timeout = time.Duration(preParse.StreamTimeout) * time.Second
 	}
 
+	if preParse.KeepAlive <= 0 {
+		sta.KeepAlive = -1
+	} else {
+		sta.KeepAlive = time.Duration(preParse.KeepAlive) * time.Second
+	}
+
 	sta.RedirHost, sta.RedirPort, err = parseRedirAddr(preParse.RedirAddr)
 	if err != nil {
 		return fmt.Errorf("unable to parse RedirAddr: %v", err)