94 lines
2 KiB
Go
94 lines
2 KiB
Go
package graceful
|
|
|
|
import (
|
|
"crypto/tls"
|
|
"errors"
|
|
"net"
|
|
"net/http"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
// TLSListener is a TLS wrapper for net.TCPListener
|
|
type TLSListener struct {
|
|
// protects conns
|
|
sync.Mutex
|
|
|
|
*net.TCPListener
|
|
|
|
config *tls.Config
|
|
conns map[net.Conn]struct{}
|
|
keepAlive time.Duration
|
|
}
|
|
|
|
// TLSConn is a net.Conn wrapper
|
|
type TLSConn struct {
|
|
net.Conn
|
|
*TLSListener
|
|
}
|
|
|
|
// Close wraps the net.Conn.Close method and removes the conn from the listener map
|
|
func (c *TLSConn) Close() error {
|
|
c.RemoveConn(c.Conn)
|
|
return c.Conn.Close()
|
|
}
|
|
|
|
// NewTLS wraps TLSListener around net.TCPListener
|
|
func NewTLS(laddr string, config *tls.Config, keepAlive time.Duration) (*TLSListener, error) {
|
|
if config == nil || len(config.Certificates) == 0 {
|
|
return nil, errors.New("NewTLS: no certificates in configuration")
|
|
}
|
|
l, err := net.Listen(network, laddr)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
ltcp, ok := l.(*net.TCPListener)
|
|
if !ok {
|
|
return nil, ErrInterfaceConversion
|
|
}
|
|
return &TLSListener{
|
|
TCPListener: ltcp,
|
|
config: config,
|
|
conns: make(map[net.Conn]struct{}),
|
|
keepAlive: keepAlive,
|
|
}, nil
|
|
}
|
|
|
|
// Accept wraps the net.TCPListener.AcceptTCP method
|
|
func (l *TLSListener) Accept() (net.Conn, error) {
|
|
conn, err := l.AcceptTCP()
|
|
if err != nil {
|
|
return conn, err
|
|
}
|
|
if l.keepAlive > 0 {
|
|
conn.SetKeepAlive(true)
|
|
conn.SetKeepAlivePeriod(l.keepAlive)
|
|
}
|
|
tlsconn := &TLSConn{tls.Server(conn, l.config), l}
|
|
l.Lock()
|
|
defer l.Unlock()
|
|
l.conns[tlsconn] = struct{}{}
|
|
return tlsconn, nil
|
|
}
|
|
|
|
// TLSConnState implements HTTP connection hijacking
|
|
func (l *TLSListener) TLSConnState(c net.Conn, cs http.ConnState) {
|
|
if cs == http.StateClosed || cs == http.StateHijacked {
|
|
l.RemoveConn(c)
|
|
}
|
|
}
|
|
|
|
// RemoveConn removes a connection from the map
|
|
func (l *TLSListener) RemoveConn(c net.Conn) {
|
|
l.Lock()
|
|
defer l.Unlock()
|
|
delete(l.conns, c)
|
|
}
|
|
|
|
// Stop stops the listener and closes all open connections
|
|
func (l *TLSListener) Stop() {
|
|
l.Close()
|
|
for c := range l.conns {
|
|
c.Close()
|
|
}
|
|
}
|