init: 自 zomaintain/backend/rdplib 平移独立成库; module path 从上游 github.com/nakagami/grdp 改为 git.zeroonesoft.cn/golib/rdplib
This commit is contained in:
@@ -0,0 +1,293 @@
|
||||
package tpkt
|
||||
|
||||
import (
|
||||
"encoding/binary"
|
||||
"fmt"
|
||||
"io"
|
||||
"log/slog"
|
||||
"runtime/debug"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"git.zeroonesoft.cn/golib/rdplib/core"
|
||||
"git.zeroonesoft.cn/golib/rdplib/emission"
|
||||
"git.zeroonesoft.cn/golib/rdplib/protocol/nla"
|
||||
)
|
||||
|
||||
var writePool = sync.Pool{
|
||||
New: func() any { return make([]byte, 0, 4096) },
|
||||
}
|
||||
|
||||
// readBufPool reuses packet body buffers in readLoop.
|
||||
// Typical RDP packets are well under 4 KiB; the pool avoids a heap
|
||||
// allocation for every incoming packet (~60/s during active sessions).
|
||||
// Buffers larger than maxPooledReadBuf are not pooled to avoid keeping
|
||||
// large slices alive in the pool between bursts.
|
||||
const maxPooledReadBuf = 32 * 1024
|
||||
|
||||
var readBufPool = sync.Pool{
|
||||
New: func() any { return make([]byte, 0, 4096) },
|
||||
}
|
||||
|
||||
func acquireReadBuf(n int) []byte {
|
||||
b := readBufPool.Get().([]byte)
|
||||
if cap(b) >= n {
|
||||
return b[:n]
|
||||
}
|
||||
readBufPool.Put(b[:0])
|
||||
return make([]byte, n)
|
||||
}
|
||||
|
||||
func releaseReadBuf(b []byte) {
|
||||
if cap(b) <= maxPooledReadBuf {
|
||||
readBufPool.Put(b[:0])
|
||||
}
|
||||
}
|
||||
|
||||
// take idea from https://github.com/Madnikulin50/gordp
|
||||
|
||||
/**
|
||||
* Type of tpkt packet
|
||||
* Fastpath is use to shortcut RDP stack
|
||||
* @see http://msdn.microsoft.com/en-us/library/cc240621.aspx
|
||||
* @see http://msdn.microsoft.com/en-us/library/cc240589.aspx
|
||||
*/
|
||||
const (
|
||||
FASTPATH_ACTION_FASTPATH = 0x0
|
||||
FASTPATH_ACTION_X224 = 0x3
|
||||
)
|
||||
|
||||
/**
|
||||
* TPKT layer of rdp stack
|
||||
*/
|
||||
type TPKT struct {
|
||||
emission.Emitter
|
||||
Conn *core.SocketLayer
|
||||
ntlm *nla.NTLMv2
|
||||
fastPathListener core.FastPathListener
|
||||
ntlmSec *nla.NTLMv2Security
|
||||
}
|
||||
|
||||
func New(s *core.SocketLayer, ntlm *nla.NTLMv2) *TPKT {
|
||||
t := &TPKT{
|
||||
Emitter: *emission.NewEmitter(),
|
||||
Conn: s,
|
||||
ntlm: ntlm,
|
||||
}
|
||||
go t.readLoop()
|
||||
return t
|
||||
}
|
||||
|
||||
// readLoop is the single goroutine that reads all incoming TPKT/FastPath packets.
|
||||
// It replaces the previous callback-chain pattern (StartReadBytes → recvHeader →
|
||||
// StartReadBytes → recvExtendedHeader → …) which spawned a new goroutine for each
|
||||
// individual read. By using a single blocking loop with io.ReadFull, we eliminate
|
||||
// goroutine creation/destruction overhead on the hot receive path.
|
||||
func (t *TPKT) readLoop() {
|
||||
var hdr [2]byte
|
||||
for {
|
||||
if _, err := io.ReadFull(t.Conn, hdr[:]); err != nil {
|
||||
t.Emit("error", err)
|
||||
return
|
||||
}
|
||||
|
||||
version := hdr[0]
|
||||
if version == FASTPATH_ACTION_X224 {
|
||||
// TPKT packet: 4-byte header total (version, reserved, length-hi, length-lo)
|
||||
var extHdr [2]byte
|
||||
if _, err := io.ReadFull(t.Conn, extHdr[:]); err != nil {
|
||||
t.Emit("error", err)
|
||||
return
|
||||
}
|
||||
size := binary.BigEndian.Uint16(extHdr[:])
|
||||
if size < 4 {
|
||||
t.Emit("error", fmt.Errorf("TPKT: invalid packet size %d", size))
|
||||
return
|
||||
}
|
||||
body := acquireReadBuf(int(size) - 4)
|
||||
if _, err := io.ReadFull(t.Conn, body); err != nil {
|
||||
t.Emit("error", err)
|
||||
return
|
||||
}
|
||||
t.Emit("data", body)
|
||||
releaseReadBuf(body)
|
||||
} else {
|
||||
// FastPath packet: 2- or 3-byte header
|
||||
secFlag := (version >> 6) & 0x3
|
||||
length := int(hdr[1])
|
||||
slog.Debug("TPKT FastPath", "secFlag", secFlag, "length", length)
|
||||
|
||||
var packetSize int
|
||||
if length&0x80 != 0 {
|
||||
// Extended 3-byte header: high 7 bits from hdr[1], low 8 from next byte
|
||||
var extByte [1]byte
|
||||
if _, err := io.ReadFull(t.Conn, extByte[:]); err != nil {
|
||||
slog.Error("TPKT recvExtendedFastPathHeader", "err", err)
|
||||
return
|
||||
}
|
||||
leftPart := length & ^0x80
|
||||
packetSize = (leftPart<<8) + int(extByte[0]) - 3
|
||||
} else {
|
||||
packetSize = length - 2
|
||||
}
|
||||
|
||||
if packetSize < 0 {
|
||||
t.Emit("error", fmt.Errorf("TPKT FastPath: invalid packet size %d", packetSize))
|
||||
return
|
||||
}
|
||||
body := acquireReadBuf(packetSize)
|
||||
if _, err := io.ReadFull(t.Conn, body); err != nil {
|
||||
slog.Debug("TPKT recvFastPath error", "err", err)
|
||||
return
|
||||
}
|
||||
t.recvFastPathSafe(secFlag, body)
|
||||
releaseReadBuf(body)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (t *TPKT) StartTLS() error {
|
||||
return t.Conn.StartTLS()
|
||||
}
|
||||
|
||||
// recvFastPathSafe 隔离 fast-path 解码 panic：解码器遇到意外数据时记录
|
||||
// 堆栈并上报错误横幅，而不是让整个 Go 运行时崩溃（worker 只会显示
|
||||
// "Go program has already exited"，毫无诊断信息）。
|
||||
func (t *TPKT) recvFastPathSafe(secFlag byte, body []byte) {
|
||||
defer func() {
|
||||
if r := recover(); r != nil {
|
||||
st := strings.Split(string(debug.Stack()), "\n")
|
||||
if len(st) > 14 {
|
||||
st = st[:14]
|
||||
}
|
||||
slog.Error("RecvFastPath panic", "recover", r, "stack", strings.Join(st, "\n"))
|
||||
t.Emit("error", fmt.Errorf("fast-path panic: %v\n%s", r, strings.Join(st, "\n")))
|
||||
}
|
||||
}()
|
||||
t.fastPathListener.RecvFastPath(secFlag, body)
|
||||
}
|
||||
|
||||
func (t *TPKT) StartNLA() error {
|
||||
// Set a deadline for the entire NLA handshake (TLS + NTLM auth)
|
||||
// to prevent hanging indefinitely when the server is slow to respond.
|
||||
t.Conn.SetDeadline(time.Now().Add(30 * time.Second))
|
||||
defer t.Conn.SetDeadline(time.Time{}) // clear deadline after NLA completes
|
||||
|
||||
slog.Debug("StartNLA: TLS handshake begin")
|
||||
err := t.StartTLS()
|
||||
if err != nil {
|
||||
slog.Error("StartNLA", "start tls failed", err)
|
||||
return err
|
||||
}
|
||||
slog.Debug("StartNLA: TLS handshake complete")
|
||||
req := nla.EncodeDERTRequest([]nla.Message{t.ntlm.GetNegotiateMessage()}, nil, nil)
|
||||
slog.Debug("StartNLA send", "req", core.Hex(req), "len", len(req))
|
||||
_, err = t.Conn.Write(req)
|
||||
if err != nil {
|
||||
slog.Error("send NegotiateMessage", "err", err)
|
||||
return err
|
||||
}
|
||||
|
||||
resp := make([]byte, 1024)
|
||||
n, err := t.Conn.Read(resp)
|
||||
slog.Debug("StartNLA recv", "n", n, "err", err)
|
||||
if err != nil {
|
||||
return fmt.Errorf("read %s", err)
|
||||
} else {
|
||||
slog.Debug("StartNLA Read success")
|
||||
}
|
||||
return t.recvChallenge(resp[:n])
|
||||
}
|
||||
|
||||
func (t *TPKT) recvChallenge(data []byte) error {
|
||||
slog.Debug("recvChallenge", "data", core.Hex(data))
|
||||
tsreq, err := nla.DecodeDERTRequest(data)
|
||||
if err != nil {
|
||||
slog.Debug("DecodeDERTRequest", "err", err)
|
||||
return err
|
||||
}
|
||||
slog.Debug("recvChallenge", "tsreq", tsreq)
|
||||
// get pubkey
|
||||
pubkey, err := t.Conn.TlsPubKey()
|
||||
slog.Debug("recvChallenge", "pubkey", core.Hex(pubkey))
|
||||
|
||||
authMsg, ntlmSec := t.ntlm.GetAuthenticateMessage(tsreq.NegoTokens[0].Data)
|
||||
t.ntlmSec = ntlmSec
|
||||
|
||||
encryptPubkey := ntlmSec.GssEncrypt(pubkey)
|
||||
req := nla.EncodeDERTRequest([]nla.Message{authMsg}, nil, encryptPubkey)
|
||||
slog.Debug("recvChallenge", "send", core.Hex(req), "len", len(req))
|
||||
_, err = t.Conn.Write(req)
|
||||
if err != nil {
|
||||
slog.Error("send AuthenticateMessage", "err", err)
|
||||
return err
|
||||
}
|
||||
|
||||
slog.Debug("recvChallenge read challenge start")
|
||||
resp := make([]byte, 1024)
|
||||
n, err := t.Conn.Read(resp)
|
||||
if err != nil {
|
||||
slog.Error("recvChallenge", "err", err)
|
||||
return fmt.Errorf("read %s", err)
|
||||
}
|
||||
|
||||
return t.recvPubKeyInc(resp[:n])
|
||||
}
|
||||
|
||||
func (t *TPKT) recvPubKeyInc(data []byte) error {
|
||||
slog.Debug("recvPubKeyInc", "data", core.Hex(data), "len", len(data))
|
||||
tsreq, err := nla.DecodeDERTRequest(data)
|
||||
if err != nil {
|
||||
slog.Debug("DecodeDERTRequest", "err", err)
|
||||
return err
|
||||
}
|
||||
slog.Debug("PubKeyAuth", "key", core.Hex(tsreq.PubKeyAuth))
|
||||
//ignore
|
||||
pubkey := t.ntlmSec.GssDecrypt([]byte(tsreq.PubKeyAuth))
|
||||
slog.Debug("GssDecrypy", "pubkey", core.Hex(pubkey))
|
||||
domain, username, password := t.ntlm.GetEncodedCredentials()
|
||||
credentials := nla.EncodeDERTCredentials(domain, username, password)
|
||||
authInfo := t.ntlmSec.GssEncrypt(credentials)
|
||||
req := nla.EncodeDERTRequest(nil, authInfo, nil)
|
||||
_, err = t.Conn.Write(req)
|
||||
if err != nil {
|
||||
slog.Debug("send AuthenticateMessage", "err", err)
|
||||
return err
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (t *TPKT) Read(b []byte) (n int, err error) {
|
||||
return t.Conn.Read(b)
|
||||
}
|
||||
|
||||
func (t *TPKT) Write(data []byte) (n int, err error) {
|
||||
buf := writePool.Get().([]byte)
|
||||
size := uint16(len(data) + 4)
|
||||
buf = append(buf[:0], FASTPATH_ACTION_X224, 0, byte(size>>8), byte(size))
|
||||
buf = append(buf, data...)
|
||||
n, err = t.Conn.Write(buf)
|
||||
writePool.Put(buf[:0])
|
||||
return
|
||||
}
|
||||
|
||||
func (t *TPKT) Close() error {
|
||||
return t.Conn.Close()
|
||||
}
|
||||
|
||||
func (t *TPKT) SetFastPathListener(f core.FastPathListener) {
|
||||
t.fastPathListener = f
|
||||
}
|
||||
|
||||
func (t *TPKT) SendFastPath(secFlag byte, data []byte) (n int, err error) {
|
||||
buf := writePool.Get().([]byte)
|
||||
hdr := uint16(len(data)+3) | 0x8000
|
||||
buf = append(buf[:0], FASTPATH_ACTION_FASTPATH|((secFlag&0x3)<<6), byte(hdr>>8), byte(hdr))
|
||||
buf = append(buf, data...)
|
||||
n, err = t.Conn.Write(buf)
|
||||
writePool.Put(buf[:0])
|
||||
return
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user