add tls support to stratum (closes #2)

etclabscore · Jan 21, 2021 · 1869677 · 1869677
1 parent a9ed9a7
commit 1869677
Show file tree

Hide file tree

Showing 4 changed files with 34 additions and 16 deletions.
diff --git a/config.example.json b/config.example.json
@@ -22,7 +22,10 @@
 			"enabled": true,
 			"listen": "0.0.0.0:8008",
 			"timeout": "120s",
-			"maxConn": 8192
+			"maxConn": 8192,
+			"tls": false,
+      "certFile": "/path/to/cert.pem",
+      "keyFile": "/path/to/key.pem"
 		},
 
 		"policy": {

diff --git a/proxy/config.go b/proxy/config.go
@@ -51,10 +51,13 @@ type Proxy struct {
 }
 
 type Stratum struct {
-	Enabled bool   `json:"enabled"`
-	Listen  string `json:"listen"`
-	Timeout string `json:"timeout"`
-	MaxConn int    `json:"maxConn"`
+	Enabled  bool   `json:"enabled"`
+	Listen   string `json:"listen"`
+	Timeout  string `json:"timeout"`
+	MaxConn  int    `json:"maxConn"`
+	TLS      bool   `json:"tls`
+	CertFile string `json:"certFile`
+	KeyFile  string `json:"keyFile`
 }
 
 type StratumNiceHash struct {

diff --git a/proxy/proxy.go b/proxy/proxy.go
@@ -49,7 +49,7 @@ type Session struct {
 
 	// Stratum
 	sync.Mutex
-	conn           *net.TCPConn
+	conn           net.Conn
 	login          string
 	subscriptionID string
 	JobDeatils     jobDetails

diff --git a/proxy/stratum.go b/proxy/stratum.go
@@ -2,6 +2,7 @@ package proxy
 
 import (
 	"bufio"
+	"crypto/tls"
 	"encoding/json"
 	"errors"
 	"io"
@@ -17,14 +18,25 @@ const (
 )
 
 func (s *ProxyServer) ListenTCP() {
-	timeout := util.MustParseDuration(s.config.Proxy.Stratum.Timeout)
-	s.timeout = timeout
-
-	addr, err := net.ResolveTCPAddr("tcp4", s.config.Proxy.Stratum.Listen)
-	if err != nil {
-		log.Fatalf("Error: %v", err)
+	s.timeout = util.MustParseDuration(s.config.Proxy.Stratum.Timeout)
+
+	var err error
+	var server net.Listener
+	setKeepAlive := func(net.Conn) {}
+	if s.config.Proxy.Stratum.TLS {
+		var cert tls.Certificate
+		cert, err = tls.LoadX509KeyPair(s.config.Proxy.Stratum.CertFile, s.config.Proxy.Stratum.KeyFile)
+		if err != nil {
+			log.Fatalln("Error loading certificate:", err)
+		}
+		tlsCfg := &tls.Config{Certificates: []tls.Certificate{cert}}
+		server, err = tls.Listen("tcp", s.config.Proxy.Stratum.Listen, tlsCfg)
+	} else {
+		server, err = net.Listen("tcp", s.config.Proxy.Stratum.Listen)
+		setKeepAlive = func(conn net.Conn) {
+			conn.(*net.TCPConn).SetKeepAlive(true)
+		}
 	}
-	server, err := net.ListenTCP("tcp4", addr)
 	if err != nil {
 		log.Fatalf("Error: %v", err)
 	}
@@ -35,11 +47,11 @@ func (s *ProxyServer) ListenTCP() {
 	n := 0
 
 	for {
-		conn, err := server.AcceptTCP()
+		conn, err := server.Accept()
 		if err != nil {
 			continue
 		}
-		conn.SetKeepAlive(true)
+		setKeepAlive(conn)
 
 		ip, _, _ := net.SplitHostPort(conn.RemoteAddr().String())
 
@@ -168,7 +180,7 @@ func (cs *Session) sendTCPError(id json.RawMessage, reply *ErrorReply) error {
 	return errors.New(reply.Message)
 }
 
-func (self *ProxyServer) setDeadline(conn *net.TCPConn) {
+func (self *ProxyServer) setDeadline(conn net.Conn) {
 	conn.SetDeadline(time.Now().Add(self.timeout))
 }