added TLS transport

This commit is contained in:
2023-05-18 11:12:16 +02:00
parent 4716c25cf6
commit aed740c933
2 changed files with 291 additions and 218 deletions

View File

@@ -6,6 +6,7 @@ package dsrpc
import ( import (
"context" "context"
"crypto/tls"
"errors" "errors"
"fmt" "fmt"
"io" "io"
@@ -32,6 +33,18 @@ func Put(ctx context.Context, address string, method string, reader io.Reader, b
return ConnPut(ctx, conn, method, reader, binSize, param, result, auth) return ConnPut(ctx, conn, method, reader, binSize, param, result, auth)
} }
func PutTLS(ctx context.Context, tlsConfig *tls.Config, address string, method string, reader io.Reader, binSize int64, param, result any, auth *Auth) error {
var err error
conn, err := tls.Dial("tcp", address, tlsConfig)
if err != nil {
return err
}
defer conn.Close()
return ConnPut(ctx, conn, method, reader, binSize, param, result, auth)
}
func ConnPut(ctx context.Context, conn net.Conn, method string, reader io.Reader, binSize int64, param, result any, auth *Auth) error { func ConnPut(ctx context.Context, conn net.Conn, method string, reader io.Reader, binSize int64, param, result any, auth *Auth) error {
var err error var err error
content := CreateContent(conn) content := CreateContent(conn)
@@ -99,6 +112,18 @@ func Get(ctx context.Context, address string, method string, writer io.Writer, p
return ConnGet(ctx, conn, method, writer, param, result, auth) return ConnGet(ctx, conn, method, writer, param, result, auth)
} }
func GetTLS(ctx context.Context, tlsConfig *tls.Config, address string, method string, writer io.Writer, param, result any, auth *Auth) error {
var err error
conn, err := tls.Dial("tcp", address, tlsConfig)
if err != nil {
return err
}
defer conn.Close()
return ConnGet(ctx, conn, method, writer, param, result, auth)
}
func ConnGet(ctx context.Context, conn net.Conn, method string, writer io.Writer, param, result any, auth *Auth) error { func ConnGet(ctx context.Context, conn net.Conn, method string, writer io.Writer, param, result any, auth *Auth) error {
var err error var err error
@@ -161,6 +186,22 @@ func Exec(ctx context.Context, address, method string, param any, result any, au
return err return err
} }
func ExecTLS(ctx context.Context, tlsConfig *tls.Config, address, method string, param any, result any, auth *Auth) error {
var err error
conn, err := tls.Dial("tcp", address, tlsConfig)
if err != nil {
return err
}
defer conn.Close()
err = ConnExec(ctx, conn, method, param, result, auth)
if err != nil {
return err
}
return err
}
func ConnExec(ctx context.Context, conn net.Conn, method string, param any, result any, auth *Auth) error { func ConnExec(ctx context.Context, conn net.Conn, method string, param any, result any, auth *Auth) error {
var err error var err error

468
server.go
View File

@@ -5,303 +5,335 @@
package dsrpc package dsrpc
import ( import (
"context" "context"
"errors" "crypto/tls"
"fmt" "errors"
"io" "fmt"
"net" "io"
"sync" "net"
"time" "sync"
"time"
encoder "encoding/json" encoder "encoding/json"
) )
type HandlerFunc = func(*Content) error type HandlerFunc = func(*Content) error
type Service struct { type Service struct {
handlers map[string]HandlerFunc handlers map[string]HandlerFunc
ctx context.Context ctx context.Context
cancel context.CancelFunc cancel context.CancelFunc
wg *sync.WaitGroup wg *sync.WaitGroup
preMw []HandlerFunc preMw []HandlerFunc
postMw []HandlerFunc postMw []HandlerFunc
keepalive bool keepalive bool
kaTime time.Duration kaTime time.Duration
kaMtx sync.Mutex kaMtx sync.Mutex
} }
func NewService() *Service { func NewService() *Service {
rdrpc := &Service{} rdrpc := &Service{}
rdrpc.handlers = make(map[string]HandlerFunc) rdrpc.handlers = make(map[string]HandlerFunc)
ctx, cancel := context.WithCancel(context.Background()) ctx, cancel := context.WithCancel(context.Background())
rdrpc.ctx = ctx rdrpc.ctx = ctx
rdrpc.cancel = cancel rdrpc.cancel = cancel
var wg sync.WaitGroup var wg sync.WaitGroup
rdrpc.wg = &wg rdrpc.wg = &wg
rdrpc.preMw = make([]HandlerFunc, 0) rdrpc.preMw = make([]HandlerFunc, 0)
rdrpc.postMw = make([]HandlerFunc, 0) rdrpc.postMw = make([]HandlerFunc, 0)
return rdrpc return rdrpc
} }
func (svc *Service) PreMiddleware(mw HandlerFunc) { func (svc *Service) PreMiddleware(mw HandlerFunc) {
svc.preMw = append(svc.preMw, mw) svc.preMw = append(svc.preMw, mw)
} }
func (svc *Service) PostMiddleware(mw HandlerFunc) { func (svc *Service) PostMiddleware(mw HandlerFunc) {
svc.postMw = append(svc.postMw, mw) svc.postMw = append(svc.postMw, mw)
} }
func (svc *Service) Handle(method string, handler HandlerFunc) { func (svc *Service) Handle(method string, handler HandlerFunc) {
svc.handlers[method] = handler svc.handlers[method] = handler
} }
func (svc *Service) SetKeepAlive(flag bool) { func (svc *Service) SetKeepAlive(flag bool) {
svc.kaMtx.Lock() svc.kaMtx.Lock()
defer svc.kaMtx.Unlock() defer svc.kaMtx.Unlock()
svc.keepalive = true svc.keepalive = true
} }
func (svc *Service) SetKeepAlivePeriod(interval time.Duration) { func (svc *Service) SetKeepAlivePeriod(interval time.Duration) {
svc.kaMtx.Lock() svc.kaMtx.Lock()
defer svc.kaMtx.Unlock() defer svc.kaMtx.Unlock()
svc.kaTime = interval svc.kaTime = interval
} }
func (svc *Service) Listen(address string) error { func (svc *Service) Listen(address string) error {
var err error var err error
logInfo("server listen:", address) logInfo("server listen:", address)
addr, err := net.ResolveTCPAddr("tcp", address) addr, err := net.ResolveTCPAddr("tcp", address)
if err != nil { if err != nil {
err = fmt.Errorf("unable to resolve adddress: %s", err) err = fmt.Errorf("unable to resolve adddress: %s", err)
return err return err
} }
listener, err := net.ListenTCP("tcp", addr) listener, err := net.ListenTCP("tcp", addr)
if err != nil { if err != nil {
err = fmt.Errorf("unable to start listener: %s", err) err = fmt.Errorf("unable to start listener: %s", err)
return err return err
} }
for { for {
conn, err := listener.AcceptTCP() conn, err := listener.AcceptTCP()
if err != nil { if err != nil {
logError("conn accept err:", err) logError("conn accept err:", err)
} }
select { select {
case <-svc.ctx.Done(): case <-svc.ctx.Done():
return err return err
default: default:
} }
svc.wg.Add(1) svc.wg.Add(1)
go svc.handleConn(conn, svc.wg) go svc.handleTCPConn(conn, svc.wg)
} }
return err return err
}
func (svc *Service) ListenTLS(address string, tlsConfig *tls.Config) error {
var err error
logInfo("server listen:", address)
listener, err := tls.Listen("tcp", address, tlsConfig)
if err != nil {
err = fmt.Errorf("unable to start listener: %s", err)
return err
}
for {
conn, err := listener.Accept()
if err != nil {
logError("conn accept err:", err)
}
select {
case <-svc.ctx.Done():
return err
default:
}
svc.wg.Add(1)
go svc.handleConn(conn, svc.wg)
}
return err
} }
func notFound(content *Content) error { func notFound(content *Content) error {
execErr := errors.New("method not found") execErr := errors.New("method not found")
err := content.SendError(execErr) err := content.SendError(execErr)
return err return err
} }
func (svc *Service) Stop() error { func (svc *Service) Stop() error {
var err error var err error
// Disable new connection // Disable new connection
logInfo("cancel rpc accept loop") logInfo("cancel rpc accept loop")
svc.cancel() svc.cancel()
// Wait handlers // Wait handlers
logInfo("wait rpc handlers") logInfo("wait rpc handlers")
svc.wg.Wait() svc.wg.Wait()
return err return err
} }
func (svc *Service) handleConn(conn *net.TCPConn, wg *sync.WaitGroup) { func (svc *Service) handleTCPConn(conn *net.TCPConn, wg *sync.WaitGroup) {
var err error var err error
if svc.keepalive {
err = conn.SetKeepAlive(true)
if err != nil {
err = fmt.Errorf("unable to set keepalive: %s", err)
return
}
if svc.kaTime > 0 {
err = conn.SetKeepAlivePeriod(svc.kaTime)
if err != nil {
err = fmt.Errorf("unable to set keepalive period: %s", err)
return
}
}
}
svc.handleConn(conn, wg)
}
if svc.keepalive { func (svc *Service) handleConn(conn net.Conn, wg *sync.WaitGroup) {
err = conn.SetKeepAlive(true) var err error
if err != nil {
err = fmt.Errorf("unable to set keepalive: %s", err)
return
}
if svc.kaTime > 0 {
err = conn.SetKeepAlivePeriod(svc.kaTime)
if err != nil {
err = fmt.Errorf("unable to set keepalive period: %s", err)
return
}
}
}
content := CreateContent(conn)
remoteAddr := conn.RemoteAddr().String() content := CreateContent(conn)
remoteHost, _, _ := net.SplitHostPort(remoteAddr)
content.remoteHost = remoteHost
content.binReader = conn remoteAddr := conn.RemoteAddr().String()
content.binWriter = io.Discard remoteHost, _, _ := net.SplitHostPort(remoteAddr)
content.remoteHost = remoteHost
exitFunc := func() { content.binReader = conn
conn.Close() content.binWriter = io.Discard
wg.Done()
if err != nil {
logError("conn handler err:", err)
}
}
defer exitFunc()
recovFunc := func() { exitFunc := func() {
panicMsg := recover() conn.Close()
if panicMsg != nil { wg.Done()
logError("handler panic message:", panicMsg) if err != nil {
} logError("conn handler err:", err)
} }
defer recovFunc() }
defer exitFunc()
err = content.ReadRequest() recovFunc := func() {
if err != nil { panicMsg := recover()
err = err if panicMsg != nil {
return logError("handler panic message:", panicMsg)
} }
}
defer recovFunc()
err = content.BindMethod() err = content.ReadRequest()
if err != nil { if err != nil {
err = err err = err
return return
} }
for _, mw := range svc.preMw {
err = mw(content) err = content.BindMethod()
if err != nil { if err != nil {
err = err err = err
return return
} }
} for _, mw := range svc.preMw {
err = svc.Route(content) err = mw(content)
if err != nil { if err != nil {
err = err err = err
return return
} }
for _, mw := range svc.postMw { }
err = mw(content) err = svc.Route(content)
if err != nil { if err != nil {
err = err err = err
return return
} }
} for _, mw := range svc.postMw {
return err = mw(content)
if err != nil {
err = err
return
}
}
return
} }
func (svc *Service) Route(content *Content) error { func (svc *Service) Route(content *Content) error {
handler, ok := svc.handlers[content.reqBlock.Method] handler, ok := svc.handlers[content.reqBlock.Method]
if ok { if ok {
return handler(content) return handler(content)
} }
return notFound(content) return notFound(content)
} }
func (content *Content) ReadRequest() error { func (content *Content) ReadRequest() error {
var err error var err error
content.reqPacket.header, err = ReadBytes(content.sockReader, headerSize) content.reqPacket.header, err = ReadBytes(content.sockReader, headerSize)
if err != nil { if err != nil {
return err return err
} }
content.reqHeader, err = UnpackHeader(content.reqPacket.header) content.reqHeader, err = UnpackHeader(content.reqPacket.header)
if err != nil { if err != nil {
return err return err
} }
rpcSize := content.reqHeader.rpcSize rpcSize := content.reqHeader.rpcSize
content.reqPacket.rcpPayload, err = ReadBytes(content.sockReader, rpcSize) content.reqPacket.rcpPayload, err = ReadBytes(content.sockReader, rpcSize)
if err != nil { if err != nil {
return err return err
} }
return err return err
} }
func (content *Content) BinWriter() io.Writer { func (content *Content) BinWriter() io.Writer {
return content.sockWriter return content.sockWriter
} }
func (content *Content) BinReader() io.Reader { func (content *Content) BinReader() io.Reader {
return content.sockReader return content.sockReader
} }
func (content *Content) BinSize() int64 { func (content *Content) BinSize() int64 {
return content.reqHeader.binSize return content.reqHeader.binSize
} }
func (content *Content) ReadBin(ctx context.Context, writer io.Writer) error { func (content *Content) ReadBin(ctx context.Context, writer io.Writer) error {
var err error var err error
_, err = CopyBytes(ctx, content.sockReader, writer, content.reqHeader.binSize) _, err = CopyBytes(ctx, content.sockReader, writer, content.reqHeader.binSize)
return err return err
} }
func (content *Content) BindMethod() error { func (content *Content) BindMethod() error {
var err error var err error
err = encoder.Unmarshal(content.reqPacket.rcpPayload, content.reqBlock) err = encoder.Unmarshal(content.reqPacket.rcpPayload, content.reqBlock)
return err return err
} }
func (content *Content) BindParams(params any) error { func (content *Content) BindParams(params any) error {
var err error var err error
content.reqBlock.Params = params content.reqBlock.Params = params
err = encoder.Unmarshal(content.reqPacket.rcpPayload, content.reqBlock) err = encoder.Unmarshal(content.reqPacket.rcpPayload, content.reqBlock)
if err != nil { if err != nil {
return err return err
} }
return err return err
} }
func (content *Content) SendResult(result any, binSize int64) error { func (content *Content) SendResult(result any, binSize int64) error {
var err error var err error
content.resBlock.Result = result content.resBlock.Result = result
content.resPacket.rcpPayload, err = content.resBlock.Pack() content.resPacket.rcpPayload, err = content.resBlock.Pack()
if err != nil { if err != nil {
return err return err
} }
content.resHeader.rpcSize = int64(len(content.resPacket.rcpPayload)) content.resHeader.rpcSize = int64(len(content.resPacket.rcpPayload))
content.resHeader.binSize = binSize content.resHeader.binSize = binSize
content.resPacket.header, err = content.resHeader.Pack() content.resPacket.header, err = content.resHeader.Pack()
if err != nil { if err != nil {
return err return err
} }
_, err = content.sockWriter.Write(content.resPacket.header) _, err = content.sockWriter.Write(content.resPacket.header)
if err != nil { if err != nil {
return err return err
} }
_, err = content.sockWriter.Write(content.resPacket.rcpPayload) _, err = content.sockWriter.Write(content.resPacket.rcpPayload)
if err != nil { if err != nil {
return err return err
} }
return err return err
} }
func (content *Content) SendError(execErr error) error { func (content *Content) SendError(execErr error) error {
var err error var err error
content.resBlock.Error = execErr.Error() content.resBlock.Error = execErr.Error()
content.resBlock.Result = NewEmptyResult() content.resBlock.Result = NewEmptyResult()
content.resPacket.rcpPayload, err = content.resBlock.Pack() content.resPacket.rcpPayload, err = content.resBlock.Pack()
if err != nil { if err != nil {
return err return err
} }
content.resHeader.rpcSize = int64(len(content.resPacket.rcpPayload)) content.resHeader.rpcSize = int64(len(content.resPacket.rcpPayload))
content.resPacket.header, err = content.resHeader.Pack() content.resPacket.header, err = content.resHeader.Pack()
if err != nil { if err != nil {
return err return err
} }
_, err = content.sockWriter.Write(content.resPacket.header) _, err = content.sockWriter.Write(content.resPacket.header)
if err != nil { if err != nil {
return err return err
} }
_, err = content.sockWriter.Write(content.resPacket.rcpPayload) _, err = content.sockWriter.Write(content.resPacket.rcpPayload)
if err != nil { if err != nil {
return err return err
} }
return err return err
} }