88 lines
2.0 KiB
Go
88 lines
2.0 KiB
Go
|
package tls
|
||
|
|
|
||
|
|
import (
|
||
|
|
"github.com/lesismal/llib/std/crypto/tls"
|
||
|
|
"github.com/lesismal/nbio"
|
||
|
|
"github.com/lesismal/nbio/mempool"
|
||
|
|
)
|
||
|
|
|
||
|
|
// Conn .
|
||
|
|
type Conn = tls.Conn
|
||
|
|
|
||
|
|
// Config .
|
||
|
|
type Config = tls.Config
|
||
|
|
|
||
|
|
// Dial returns a net.Conn to be added to a Gopher.
|
||
|
|
func Dial(network, addr string, config *Config) (*tls.Conn, error) {
|
||
|
|
tlsConn, err := tls.Dial(network, addr, config, mempool.DefaultMemPool)
|
||
|
|
if err != nil {
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
|
||
|
|
return tlsConn, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
// WrapOpen returns an opening handler of nbio.Gopher.
|
||
|
|
func WrapOpen(tlsConfig *Config, isClient bool, h func(c *nbio.Conn, tlsConn *Conn)) func(c *nbio.Conn) {
|
||
|
|
return func(c *nbio.Conn) {
|
||
|
|
var tlsConn *tls.Conn
|
||
|
|
sesseion := c.Session()
|
||
|
|
if sesseion != nil {
|
||
|
|
tlsConn = sesseion.(*tls.Conn)
|
||
|
|
}
|
||
|
|
if tlsConn == nil && !isClient {
|
||
|
|
tlsConn = tls.NewConn(c, tlsConfig, isClient, true, mempool.DefaultMemPool)
|
||
|
|
c.SetSession(tlsConn)
|
||
|
|
}
|
||
|
|
if h != nil {
|
||
|
|
h(c, tlsConn)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// WrapClose returns an closing handler of nbio.Gopher.
|
||
|
|
func WrapClose(h func(c *nbio.Conn, tlsConn *Conn, err error)) func(c *nbio.Conn, err error) {
|
||
|
|
return func(c *nbio.Conn, err error) {
|
||
|
|
if h != nil && c != nil {
|
||
|
|
if session := c.Session(); session != nil {
|
||
|
|
if tlsConn, ok := session.(*Conn); ok {
|
||
|
|
h(c, tlsConn, err)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// WrapData returns a data handler of nbio.Gopher.
|
||
|
|
func WrapData(h func(c *nbio.Conn, tlsConn *Conn, data []byte), args ...interface{}) func(c *nbio.Conn, data []byte) {
|
||
|
|
getBuffer := func() []byte {
|
||
|
|
return make([]byte, 2048)
|
||
|
|
}
|
||
|
|
if len(args) > 0 {
|
||
|
|
if bh, ok := args[0].(func() []byte); ok {
|
||
|
|
getBuffer = bh
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return func(c *nbio.Conn, data []byte) {
|
||
|
|
if session := c.Session(); session != nil {
|
||
|
|
if tlsConn, ok := session.(*Conn); ok {
|
||
|
|
tlsConn.Append(data)
|
||
|
|
buffer := getBuffer()
|
||
|
|
for {
|
||
|
|
n, err := tlsConn.Read(buffer)
|
||
|
|
if err != nil {
|
||
|
|
c.Close()
|
||
|
|
return
|
||
|
|
}
|
||
|
|
if h != nil && n > 0 {
|
||
|
|
h(c, tlsConn, buffer[:n])
|
||
|
|
}
|
||
|
|
if n == 0 {
|
||
|
|
return
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|