[ME-4546] Clean-Up Fork, Add Custom B0 Proxy Support

buf.go

@@ -45,18 +45,22 @@ type tdsBuffer struct {
 	wpos        int
 	wPacketSeq  byte
 	wPacketType packetType
+	wSpid       uint16
 
 	// Read fields.
 	rbuf        []byte
 	rpos        int
 	rsize       int
 	final       bool
 	rPacketType packetType
+	rSpid       uint16
 
 	// afterFirst is assigned to right after tdsBuffer is created and
 	// before the first use. It is executed after the first packet is
 	// written and then removed.
 	afterFirst func()
+
+	serverConn *tdsSession
 }
 
 func newTdsBuffer(bufsize uint16, transport io.ReadWriteCloser) *tdsBuffer {
@@ -86,6 +90,7 @@ func (w *tdsBuffer) flush() (err error) {
 	// Write packet size.
 	w.wbuf[0] = byte(w.wPacketType)
 	binary.BigEndian.PutUint16(w.wbuf[2:], uint16(w.wpos))
+	binary.BigEndian.PutUint16(w.wbuf[4:], w.wSpid)
 	w.wbuf[6] = w.wPacketSeq
 
 	// Write packet into underlying transport.
@@ -169,12 +174,14 @@ func (r *tdsBuffer) readNextPacket() error {
 		PacketNo:   buf[6],
 		Pad:        buf[7],
 	}
+
 	if int(h.Size) > r.packetSize {
 		return errors.New("invalid packet size, it is longer than buffer size")
 	}
 	if headerSize > int(h.Size) {
 		return errors.New("invalid packet size, it is shorter than header size")
 	}
+
 	_, err = io.ReadFull(r.transport, r.rbuf[headerSize:h.Size])
 	//s := base64.StdEncoding.EncodeToString(r.rbuf[headerSize:h.Size])
 	//fmt.Print(s)
@@ -183,14 +190,58 @@ func (r *tdsBuffer) readNextPacket() error {
 	}
 	r.rpos = headerSize
 	r.rsize = int(h.Size)
-	r.final = h.Status != 0
+	r.final = h.Status&0x1 != 0
 	r.rPacketType = h.PacketType
+	r.rSpid = h.Spid
+
+	if r.serverConn != nil {
+		_, err := r.serverConn.buf.Write(r.rbuf[r.rpos:r.rsize])
+		if err != nil {
+			return err
+		}
+
+		if r.final {
+			r.serverConn.buf.wSpid = h.Spid
+			if err := r.serverConn.buf.FinishPacket(); err != nil {
+				return err
+			}
+		} else {
+			if err := r.serverConn.buf.flush(); err != nil {
+				return err
+			}
+		}
+	}
 	return nil
 }
 
-func (r *tdsBuffer) BeginRead() (packetType, error) {
+type beginReadConfig struct {
+	id string
+}
+
+type beginRedOption func(beginReadConfig)
+
+func withFallbackID(id string) beginRedOption { return func(brc beginReadConfig) { brc.id = id } }
+
+func (r *tdsBuffer) BeginRead(opts ...beginRedOption) (packetType, error) {
+	conf := beginReadConfig{
+		id: "UNSET",
+	}
+	for _, opt := range opts {
+		opt(conf)
+	}
+
+	// var a string
+	if r.serverConn != nil {
+		conf.id = r.serverConn.id
+		// fmt.Printf("BeginRead with serverconn: %s %d\n", r.serverConn.id, r.rPacketType)
+		r.serverConn.buf.BeginPacket(r.rPacketType, false)
+	} else {
+		// fmt.Printf("BeginRead no serverconn: %s %d\n", conf.id, r.rPacketType)
+	}
+
 	err := r.readNextPacket()
 	if err != nil {
+		// fmt.Printf("BeginRead error: %s: %s\n", conf.id, err)
 		return 0, err
 	}
 	return r.rPacketType, nil

go.sum

@@ -13,8 +13,6 @@ github.com/AzureAD/microsoft-authentication-library-for-go v1.2.2/go.mod h1:wP83
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
-github.com/golang-jwt/jwt/v5 v5.2.1 h1:OuVbFODueb089Lh128TAcimifWaLhJwVflnrgM17wHk=
-github.com/golang-jwt/jwt/v5 v5.2.1/go.mod h1:pqrtFR0X4osieyHYxtmOUWsAWrfe1Q5UVIyoH402zdk=
 github.com/golang-jwt/jwt/v5 v5.2.2 h1:Rl4B7itRWVtYIHFrSNd7vhTiz9UpLdi6gZhZ3wEeDy8=
 github.com/golang-jwt/jwt/v5 v5.2.2/go.mod h1:pqrtFR0X4osieyHYxtmOUWsAWrfe1Q5UVIyoH402zdk=
 github.com/golang-sql/civil v0.0.0-20220223132316-b832511892a9 h1:au07oEsX2xN0ktxqI+Sida1w446QrXBRJ0nee3SNZlA=

mssql.go

@@ -239,6 +239,10 @@ type Conn struct {
 	outs outputs
 }
 
+func (c *Conn) Spid() uint16 {
+	return c.sess.buf.rSpid
+}
+
 type outputs struct {
 	params       map[string]interface{}
 	returnStatus *ReturnStatus