ScuroNeko/mtg
1
0
Fork 0
mirror of https://github.com/ScuroNeko/mtg.git synced 2026-08-31 12:04:03 +03:00
Code Issues Packages Projects Releases Wiki Activity

Merge remote-tracking branch 'origin/stable' into v2

Browse Source
This commit is contained in:
9seconds
2026-03-19 18:52:35 +01:00
parent dc81f7981c cb627f2a66
commit fc72de9e39
7 changed files with 165 additions and 190 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mtglib/internal/doppel/conn.go
+27 -41
View File
@@ -20,44 +20,21 @@ type connPayload struct {
ctxCancel context.CancelCauseFunc
clock Clock
wg sync.WaitGroup
syncWriteLock sync.RWMutex
writeStream bytes.Buffer
writeCond *sync.Cond
writtenCond sync.Cond
done bool
}
func (c Conn) Write(p []byte) (int, error) {
c.p.syncWriteLock.RLock()
defer c.p.syncWriteLock.RUnlock()
if len(p) == 0 {
return 0, context.Cause(c.p.ctx)
}
c.p.writeCond.L.Lock()
c.p.writtenCond.L.Lock()
c.p.writeStream.Write(p)
c.p.writeCond.L.Unlock()
c.p.writtenCond.L.Unlock()
return len(p), context.Cause(c.p.ctx)
}
func (c Conn) SyncWrite(p []byte) (int, error) {
c.p.syncWriteLock.Lock()
defer c.p.syncWriteLock.Unlock()
c.p.writeCond.L.Lock()
// wait until buffer is exhausted
for c.p.writeStream.Len() != 0 && context.Cause(c.p.ctx) == nil {
c.p.writeCond.Wait()
}
c.p.writeStream.Write(p)
c.p.writeCond.L.Unlock()
if err := context.Cause(c.p.ctx); err != nil {
return len(p), err
}
c.p.writeCond.L.Lock()
// wait until data will be sent
for c.p.writeStream.Len() != 0 && context.Cause(c.p.ctx) == nil {
c.p.writeCond.Wait()
}
c.p.writeCond.L.Unlock()
c.p.writtenCond.Signal()
return len(p), context.Cause(c.p.ctx)
}
@@ -69,8 +46,6 @@ func (c Conn) Start() {
}
func (c Conn) start() {
defer c.p.writeCond.Broadcast()
buf := [tls.MaxRecordSize]byte{}
for {
@@ -80,25 +55,34 @@ func (c Conn) start() {
case <-c.p.clock.tick:
}
c.p.writeCond.L.Lock()
n, err := c.p.writeStream.Read(buf[:c.p.clock.stats.Size()])
c.p.writeCond.L.Unlock()
size := c.p.clock.stats.Size()
if n == 0 || err != nil {
c.p.writtenCond.L.Lock()
for c.p.writeStream.Len() == 0 && !c.p.done {
c.p.writtenCond.Wait()
}
n, _ := c.p.writeStream.Read(buf[tls.SizeHeader : tls.SizeHeader+size])
c.p.writtenCond.L.Unlock()
if n == 0 {
continue
}
if err := tls.WriteRecord(c.Conn, buf[:n]); err != nil {
if err := tls.WriteRecordInPlace(c.Conn, buf[:], n); err != nil {
c.p.ctxCancel(err)
return
}
c.p.writeCond.Signal()
}
}
func (c Conn) Stop() {
c.p.ctxCancel(nil)
c.p.writtenCond.L.Lock()
c.p.done = true
c.p.writtenCond.L.Unlock()
c.p.writtenCond.Broadcast()
c.p.wg.Wait()
}
@@ -109,7 +93,9 @@ func NewConn(ctx context.Context, conn essentials.Conn, stats *Stats) Conn {
p: &connPayload{
ctx: ctx,
ctxCancel: cancel,
writeCond: sync.NewCond(&sync.Mutex{}),
writtenCond: sync.Cond{
L: &sync.Mutex{},
},
clock: Clock{
stats: stats,
tick: make(chan struct{}),
mtglib/internal/doppel/conn_test.go
+31 -130
View File
@@ -141,6 +141,37 @@ func (suite *ConnTestSuite) TestWriteReturnsErrorAfterStop() {
suite.Error(err)
}
func (suite *ConnTestSuite) TestStopDoesNotDeadlockWhenStartIsWaiting() {
suite.connMock.
On("Write", mock.AnythingOfType("[]uint8")).
Return(0, nil).
Maybe()
for range 100 {
func() {
ctx, cancel := context.WithCancel(suite.ctx)
defer cancel()
c := NewConn(ctx, suite.connMock, &Stats{
k: 2.0,
lambda: 0.01,
})
done := make(chan struct{})
go func() {
defer close(done)
c.Stop()
}()
select {
case <-done:
case <-time.After(2 * time.Second):
suite.Fail("Stop() deadlocked: start() likely stuck in writtenCond.Wait()")
}
}()
}
}
func (suite *ConnTestSuite) TestStopOnUnderlyingWriteError() {
suite.connMock.
On("Write", mock.AnythingOfType("[]uint8")).
@@ -157,136 +188,6 @@ func (suite *ConnTestSuite) TestStopOnUnderlyingWriteError() {
}, 2*time.Second, time.Millisecond)
}
func (suite *ConnTestSuite) TestSyncWriteDataSent() {
suite.connMock.
On("Write", mock.AnythingOfType("[]uint8")).
Return(0, nil).
Maybe()
c := suite.makeConn()
defer c.Stop()
payload := []byte("sync hello")
n, err := c.SyncWrite(payload)
suite.NoError(err)
suite.Equal(len(payload), n)
// SyncWrite returns only after data is flushed to the wire.
assembled := &bytes.Buffer{}
reader := bytes.NewReader(suite.connMock.Written())
for {
header := make([]byte, tls.SizeHeader)
if _, err := io.ReadFull(reader, header); err != nil {
break
}
suite.Equal(byte(tls.TypeApplicationData), header[0])
length := binary.BigEndian.Uint16(header[tls.SizeRecordType+tls.SizeVersion:])
rec := make([]byte, length)
_, err := io.ReadFull(reader, rec)
suite.NoError(err)
assembled.Write(rec)
}
suite.Equal(payload, assembled.Bytes())
}
func (suite *ConnTestSuite) TestSyncWriteDrainsBufferFirst() {
suite.connMock.
On("Write", mock.AnythingOfType("[]uint8")).
Return(0, nil).
Maybe()
c := suite.makeConn()
defer c.Stop()
// Buffer some data via async Write.
_, err := c.Write([]byte("first"))
suite.NoError(err)
// SyncWrite must drain "first" before sending "second".
n, err := c.SyncWrite([]byte("second"))
suite.NoError(err)
suite.Equal(6, n)
// All data should be on the wire now.
assembled := &bytes.Buffer{}
reader := bytes.NewReader(suite.connMock.Written())
for {
header := make([]byte, tls.SizeHeader)
if _, err := io.ReadFull(reader, header); err != nil {
break
}
length := binary.BigEndian.Uint16(header[tls.SizeRecordType+tls.SizeVersion:])
rec := make([]byte, length)
_, err := io.ReadFull(reader, rec)
suite.NoError(err)
assembled.Write(rec)
}
suite.Equal([]byte("firstsecond"), assembled.Bytes())
}
func (suite *ConnTestSuite) TestSyncWriteBlocksAsyncWrite() {
suite.connMock.
On("Write", mock.AnythingOfType("[]uint8")).
Return(0, nil).
Maybe()
c := suite.makeConn()
defer c.Stop()
// Start SyncWrite — it holds exclusive lock.
syncDone := make(chan struct{})
go func() {
defer close(syncDone)
c.SyncWrite([]byte("exclusive")) //nolint: errcheck
}()
// Give SyncWrite time to acquire the lock.
time.Sleep(10 * time.Millisecond)
// Async Write should block until SyncWrite completes.
writeDone := make(chan struct{})
go func() {
defer close(writeDone)
c.Write([]byte("blocked")) //nolint: errcheck
}()
// SyncWrite should finish first.
<-syncDone
select {
case <-writeDone:
// Write completed after SyncWrite — correct.
case <-time.After(2 * time.Second):
suite.Fail("async Write did not unblock after SyncWrite completed")
}
}
func (suite *ConnTestSuite) TestSyncWriteReturnsErrorAfterStop() {
suite.connMock.
On("Write", mock.AnythingOfType("[]uint8")).
Return(0, nil).
Maybe()
c := suite.makeConn()
c.Stop()
time.Sleep(10 * time.Millisecond)
_, err := c.SyncWrite([]byte("too late"))
suite.Error(err)
}
func TestConn(t *testing.T) {
t.Parallel()
suite.Run(t, &ConnTestSuite{})
mtglib/internal/doppel/ganger.go
+2 -1
View File
@@ -98,7 +98,8 @@ func (g *Ganger) run() {
g.durations = append(g.durations, durations...)
if len(g.durations) > DoppelGangerMaxDurations {
g.durations = g.durations[len(g.durations)-DoppelGangerMaxDurations:]
copy(g.durations, g.durations[len(g.durations)-DoppelGangerMaxDurations:])
g.durations = g.durations[:DoppelGangerMaxDurations]
}
if len(g.durations) < MinDurationsToCalculate {
mtglib/internal/tls/conn.go
-2
View File
@@ -34,7 +34,6 @@ type Conn struct {
type connPayload struct {
readBuf bytes.Buffer
writeBuf bytes.Buffer
connBuffered *bufio.Reader
read bool
write bool
@@ -80,7 +79,6 @@ func New(conn essentials.Conn, read, write bool) Conn {
}
newConn.p.readBuf.Grow(DefaultBufferSize)
newConn.p.writeBuf.Grow(DefaultBufferSize)
return newConn
}
mtglib/internal/tls/utils.go
+15 -11
View File
@@ -29,20 +29,24 @@ func ReadRecord(r io.Reader, w io.Writer) (byte, int64, error) {
func WriteRecord(w io.Writer, payload []byte) error {
buf := [MaxRecordSize]byte{}
buf[0] = TypeApplicationData
copy(buf[SizeHeader:], payload)
bufV := buf[SizeRecordType:]
copy(bufV[:SizeVersion], TLSVersion[:])
bufS := bufV[SizeVersion:]
binary.BigEndian.PutUint16(bufS[:SizeSize], uint16(len(payload)))
bufP := buf[SizeHeader:]
if n := copy(bufP, payload); n != len(payload) {
return fmt.Errorf("copied %d bytes of payload instead of %d", n, len(payload))
return WriteRecordInPlace(w, buf[:], len(payload))
}
_, err := w.Write(buf[:SizeHeader+len(payload)])
func WriteRecordInPlace(w io.Writer, buf []byte, payloadLen int) error {
if payloadLen > MaxRecordPayloadSize {
return fmt.Errorf("payload %d exceeds max %d", payloadLen, MaxRecordPayloadSize)
}
buf[0] = TypeApplicationData
copy(buf[SizeRecordType:SizeRecordType+SizeVersion], TLSVersion[:])
binary.BigEndian.PutUint16(
buf[SizeRecordType+SizeVersion:SizeRecordType+SizeVersion+SizeSize],
uint16(payloadLen),
)
_, err := w.Write(buf[:SizeHeader+payloadLen])
return err
}
mtglib/internal/tls/utils_test.go
+78
View File
@@ -119,6 +119,84 @@ func (suite *UtilsTestSuite) TestWriteRecordPayloadTooLarge() {
suite.Error(err)
}
func (suite *UtilsTestSuite) TestWriteRecordInPlace() {
payload := []byte("hello in-place")
var buf [MaxRecordSize]byte
copy(buf[SizeHeader:], payload)
err := WriteRecordInPlace(suite.dst, buf[:], len(payload))
suite.NoError(err)
written := suite.dst.Bytes()
suite.Equal(byte(TypeApplicationData), written[0])
suite.Equal(TLSVersion[:], written[SizeRecordType:SizeRecordType+SizeVersion])
length := binary.BigEndian.Uint16(written[SizeRecordType+SizeVersion:])
suite.Equal(uint16(len(payload)), length)
suite.Equal(payload, written[SizeHeader:])
}
func (suite *UtilsTestSuite) TestWriteRecordInPlaceRoundTrip() {
payload := []byte("round trip in-place")
var buf [MaxRecordSize]byte
copy(buf[SizeHeader:], payload)
var wire bytes.Buffer
err := WriteRecordInPlace(&wire, buf[:], len(payload))
suite.NoError(err)
var recovered bytes.Buffer
recordType, length, err := ReadRecord(&wire, &recovered)
suite.NoError(err)
suite.Equal(byte(TypeApplicationData), recordType)
suite.Equal(int64(len(payload)), length)
suite.Equal(payload, recovered.Bytes())
}
func (suite *UtilsTestSuite) TestWriteRecordInPlacePayloadTooLarge() {
var buf [MaxRecordSize]byte
err := WriteRecordInPlace(suite.dst, buf[:], MaxRecordPayloadSize+1)
suite.Error(err)
}
func (suite *UtilsTestSuite) TestWriteRecordInPlacePropagatesError() {
m := &WriterMock{}
m.
On("Write", mock.AnythingOfType("[]uint8")).
Once().
Return(0, errors.New("disk full"))
var buf [MaxRecordSize]byte
copy(buf[SizeHeader:], []byte("data"))
err := WriteRecordInPlace(m, buf[:], 4)
suite.Error(err)
m.AssertExpectations(suite.T())
}
func (suite *UtilsTestSuite) TestWriteRecordInPlaceMatchesWriteRecord() {
payload := []byte("equivalence check")
var legacy bytes.Buffer
err := WriteRecord(&legacy, payload)
suite.NoError(err)
var buf [MaxRecordSize]byte
copy(buf[SizeHeader:], payload)
var inPlace bytes.Buffer
err = WriteRecordInPlace(&inPlace, buf[:], len(payload))
suite.NoError(err)
suite.Equal(legacy.Bytes(), inPlace.Bytes())
}
func TestUtils(t *testing.T) {
t.Parallel()
suite.Run(t, &UtilsTestSuite{})
mtglib/proxy.go
+8 -1
View File
@@ -259,9 +259,16 @@ func (p *Proxy) doTelegramCall(ctx *streamContext) error {
ctx: ctx,
}
telegramHost, _, err := net.SplitHostPort(foundAddr.Address)
if err != nil {
conn.Close() //nolint: errcheck
return fmt.Errorf("cannot parse telegram address %s: %w", foundAddr.Address, err)
}
p.eventStream.Send(ctx,
NewEventConnectedToDC(ctx.streamID,
conn.RemoteAddr().(*net.TCPAddr).IP, //nolint: forcetypeassert
net.ParseIP(telegramHost),
ctx.dc),
)