From f0db5e94b5f033d2bbf5879e283565e19a2e6c4a Mon Sep 17 00:00:00 2001
From: smallnest <smallnest@gmail.com>
Date: Sun, 22 Oct 2017 17:11:34 +0800
Subject: [PATCH] refactor TLS

---
 server/listener.go             |  8 +++++++-
 server/listener_all.go         | 10 +++++++---
 server/listener_windows.go     | 12 +++++++++---
 server/listener_windows_all.go | 20 ++++++++++++++------
 server/server.go               | 14 +++-----------
 5 files changed, 40 insertions(+), 24 deletions(-)

diff --git a/server/listener.go b/server/listener.go
index 2392b65..5f5dd4d 100644
--- a/server/listener.go
+++ b/server/listener.go
@@ -4,6 +4,7 @@
 package server
 
 import (
+	"crypto/tls"
 	"net"
 
 	reuseport "github.com/kavu/go_reuseport"
@@ -22,7 +23,12 @@ func (s *Server) makeListener(network, address string) (ln net.Listener, err err
 
 		ln, err = reuseport.NewReusablePortListener(network, address)
 	default: //tcp, http
-		ln, err = net.Listen(network, address)
+		if s.TLSConfig == nil {
+			ln, err = net.Listen(network, address)
+		} else {
+			ln, err = tls.Listen(network, address, s.TLSConfig)
+		}
+
 	}
 
 	return ln, err
diff --git a/server/listener_all.go b/server/listener_all.go
index c04b745..204661b 100644
--- a/server/listener_all.go
+++ b/server/listener_all.go
@@ -32,12 +32,16 @@ func (s *Server) makeListener(network, address string) (ln net.Listener, err err
 
 		ln, err = reuseport.NewReusablePortListener(network, address)
 	case "quic":
-		if s.Options == nil || s.Options["QuicConfig"] == nil {
+		if s.TLSConfig == nil {
 			return nil, errors.New("KCP BlockCrypt must be configured in server.Options")
 		}
-		ln, err = quicconn.Listen("udp", address, s.Options["QuicConfig"].(*tls.Config))
+		ln, err = quicconn.Listen("udp", address, s.TLSConfig)
 	default: //tcp, http
-		ln, err = net.Listen(network, address)
+		if s.TLSConfig == nil {
+			ln, err = net.Listen(network, address)
+		} else {
+			ln, err = tls.Listen(network, address, s.TLSConfig)
+		}
 	}
 
 	return ln, err
diff --git a/server/listener_windows.go b/server/listener_windows.go
index 709d053..22d5a98 100644
--- a/server/listener_windows.go
+++ b/server/listener_windows.go
@@ -1,15 +1,18 @@
 // +build windows
 // +build !udp
 
-package rpcx
+package server
 
 import (
+	"errors"
 	"net"
+
+	quicconn "github.com/marten-seemann/quic-conn"
 )
 
 // block can be nil if the caller wishes to skip encryption.
 // tlsConfig can be nil iff we are not using network "quic".
-func makeListener(network, address string) (ln net.Listener, err error) {
+func (s *Server) makeListener(network, address string) (ln net.Listener, err error) {
 	switch network {
 	case "reuseport":
 		if validIP4(address) {
@@ -20,7 +23,10 @@ func makeListener(network, address string) (ln net.Listener, err error) {
 
 		ln, err = net.Listen(network, address)
 	default: //tcp
-		ln, err = net.Listen(network, address)
+		if s.TLSConfig == nil {
+			return nil, errors.New("KCP BlockCrypt must be configured in server.Options")
+		}
+		ln, err = quicconn.Listen("udp", address, s.TLSConfig)
 	}
 
 	return ln, err
diff --git a/server/listener_windows_all.go b/server/listener_windows_all.go
index a04961a..6081737 100644
--- a/server/listener_windows_all.go
+++ b/server/listener_windows_all.go
@@ -1,9 +1,10 @@
 // +build windows
 // +build udp
 
-package rpcx
+package server
 
 import (
+	"crypto/tls"
 	"errors"
 	"net"
 
@@ -12,12 +13,12 @@ import (
 )
 
 // block can be nil if the caller wishes to skip encryption.
-// tlsConfig can be nil iff we are not using network "quic".
-func makeListener(network, address string) (ln net.Listener, err error) {
+// tlsConfig can be nil if we are not using network "quic".
+func (s *Server) makeListener(network, address string) (ln net.Listener, err error) {
 	switch network {
 	case "kcp":
 		if s.Options == nil || s.Options["BlockCrypt"] == nil {
-			return errors.New("KCP BlockCrypt must be configured in server.Options")
+			return nil, errors.New("KCP BlockCrypt must be configured in server.Options")
 		}
 
 		ln, err = kcp.ListenWithOptions(address, s.Options["BlockCrypt"].(kcp.BlockCrypt), 10, 3)
@@ -30,9 +31,16 @@ func makeListener(network, address string) (ln net.Listener, err error) {
 
 		ln, err = net.Listen(network, address)
 	case "quic":
-		ln, err = quicconn.Listen("udp", address, tlsConfig)
+		if s.TLSConfig == nil {
+			return nil, errors.New("KCP BlockCrypt must be configured in server.Options")
+		}
+		ln, err = quicconn.Listen("udp", address, s.TLSConfig)
 	default: //tcp
-		ln, err = net.Listen(network, address)
+		if s.TLSConfig == nil {
+			ln, err = net.Listen(network, address)
+		} else {
+			ln, err = tls.Listen(network, address, s.TLSConfig)
+		}
 	}
 
 	return ln, err
diff --git a/server/server.go b/server/server.go
index 65e3d39..d1951a7 100644
--- a/server/server.go
+++ b/server/server.go
@@ -62,7 +62,9 @@ type Server struct {
 	inShutdown int32
 	onShutdown []func()
 
-	// BlockCrypt for kcp.BlockCrypt, QUICConfig for quic TlsConfig, etc.
+	// TLSConfig for creating tls tcp connection.
+	TLSConfig *tls.Config
+	// BlockCrypt for kcp.BlockCrypt
 	Options map[string]interface{}
 	// // use for KCP
 	// KCPConfig KCPConfig
@@ -82,16 +84,6 @@ func NewServer(options map[string]interface{}) *Server {
 	}
 }
 
-// // KCPConfig is config of KCP.
-// type KCPConfig struct {
-// 	BlockCrypt kcp.BlockCrypt
-// }
-
-// // QUICConfig is config of QUIC.
-// type QUICConfig struct {
-// 	TlsConfig *tls.Config
-// }
-
 // Address returns listened address.
 func (s *Server) Address() net.Addr {
 	if s.ln == nil {