mirror of
https://github.com/ScuroNeko/mtg.git
synced 2026-08-31 11:04:02 +03:00
Add doppel and tls packages
This commit is contained in:
@@ -0,0 +1,35 @@
|
||||
package doppel
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
)
|
||||
|
||||
type Clock struct {
|
||||
stats *Stats
|
||||
tick chan struct{}
|
||||
}
|
||||
|
||||
func (c Clock) Start(ctx context.Context) {
|
||||
tickTock := time.NewTimer(c.stats.Delay())
|
||||
defer func() {
|
||||
tickTock.Stop()
|
||||
select {
|
||||
case <-tickTock.C:
|
||||
default:
|
||||
}
|
||||
}()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-tickTock.C:
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
case c.tick <- struct{}{}:
|
||||
}
|
||||
tickTock.Reset(c.stats.Delay())
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,80 @@
|
||||
package doppel
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/suite"
|
||||
)
|
||||
|
||||
type ClockTestSuite struct {
|
||||
suite.Suite
|
||||
|
||||
clock Clock
|
||||
wg sync.WaitGroup
|
||||
ctx context.Context
|
||||
ctxCancel context.CancelFunc
|
||||
}
|
||||
|
||||
func (suite *ClockTestSuite) SetupTest() {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
|
||||
suite.ctx = ctx
|
||||
suite.ctxCancel = cancel
|
||||
suite.clock = Clock{
|
||||
stats: &Stats{
|
||||
k: StatsDefaultK,
|
||||
lambda: StatsDefaultLambda,
|
||||
},
|
||||
tick: make(chan struct{}),
|
||||
}
|
||||
|
||||
suite.wg.Go(func() {
|
||||
suite.clock.Start(suite.ctx)
|
||||
})
|
||||
}
|
||||
|
||||
func (suite *ClockTestSuite) TearDownTest() {
|
||||
suite.ctxCancel()
|
||||
suite.wg.Wait()
|
||||
}
|
||||
|
||||
func (suite *ClockTestSuite) TestTicks() {
|
||||
received := 0
|
||||
|
||||
for range 3 {
|
||||
select {
|
||||
case <-suite.clock.tick:
|
||||
received++
|
||||
case <-time.After(2 * time.Second):
|
||||
suite.Fail("timed out waiting for tick")
|
||||
}
|
||||
}
|
||||
|
||||
suite.Equal(3, received)
|
||||
}
|
||||
|
||||
func (suite *ClockTestSuite) TestStopsOnCancel() {
|
||||
select {
|
||||
case <-suite.clock.tick:
|
||||
case <-time.After(2 * time.Second):
|
||||
suite.Fail("timed out waiting for first tick")
|
||||
}
|
||||
|
||||
suite.ctxCancel()
|
||||
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
|
||||
select {
|
||||
case <-suite.clock.tick:
|
||||
suite.Fail("received tick after cancel")
|
||||
default:
|
||||
}
|
||||
}
|
||||
|
||||
func TestClock(t *testing.T) {
|
||||
t.Parallel()
|
||||
suite.Run(t, &ClockTestSuite{})
|
||||
}
|
||||
@@ -0,0 +1,95 @@
|
||||
package doppel
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"sync"
|
||||
|
||||
"github.com/9seconds/mtg/v2/essentials"
|
||||
"github.com/9seconds/mtg/v2/mtglib/internal/tls"
|
||||
)
|
||||
|
||||
type Conn struct {
|
||||
essentials.Conn
|
||||
|
||||
p *connPayload
|
||||
}
|
||||
|
||||
type connPayload struct {
|
||||
ctx context.Context
|
||||
ctxCancel context.CancelCauseFunc
|
||||
clock Clock
|
||||
wg sync.WaitGroup
|
||||
writeLock sync.Mutex
|
||||
writeStream bytes.Buffer
|
||||
}
|
||||
|
||||
func (c Conn) Write(p []byte) (int, error) {
|
||||
c.p.writeLock.Lock()
|
||||
c.p.writeStream.Write(p)
|
||||
c.p.writeLock.Unlock()
|
||||
|
||||
return len(p), context.Cause(c.p.ctx)
|
||||
}
|
||||
|
||||
func (c Conn) Start() {
|
||||
c.p.wg.Go(func() {
|
||||
c.start()
|
||||
})
|
||||
}
|
||||
|
||||
func (c Conn) start() {
|
||||
buf := [tls.MaxRecordSize]byte{}
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-c.p.ctx.Done():
|
||||
return
|
||||
case <-c.p.clock.tick:
|
||||
}
|
||||
|
||||
c.p.writeLock.Lock()
|
||||
n, err := c.p.writeStream.Read(buf[:c.p.clock.stats.Size()])
|
||||
c.p.writeLock.Unlock()
|
||||
|
||||
if n == 0 || err != nil {
|
||||
continue
|
||||
}
|
||||
|
||||
if err := tls.WriteRecord(c.Conn, buf[:n]); err != nil {
|
||||
c.p.ctxCancel(err)
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (c Conn) Stop() {
|
||||
c.p.ctxCancel(nil)
|
||||
c.p.wg.Wait()
|
||||
}
|
||||
|
||||
func NewConn(ctx context.Context, conn essentials.Conn, stats *Stats) Conn {
|
||||
ctx, cancel := context.WithCancelCause(ctx)
|
||||
rv := Conn{
|
||||
Conn: conn,
|
||||
p: &connPayload{
|
||||
ctx: ctx,
|
||||
ctxCancel: cancel,
|
||||
clock: Clock{
|
||||
stats: stats,
|
||||
tick: make(chan struct{}),
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
rv.p.writeStream.Grow(tls.DefaultBufferSize)
|
||||
|
||||
rv.p.wg.Go(func() {
|
||||
rv.p.clock.Start(ctx)
|
||||
})
|
||||
rv.p.wg.Go(func() {
|
||||
rv.start()
|
||||
})
|
||||
|
||||
return rv
|
||||
}
|
||||
@@ -0,0 +1,163 @@
|
||||
package doppel
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/binary"
|
||||
"errors"
|
||||
"io"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/9seconds/mtg/v2/internal/testlib"
|
||||
"github.com/9seconds/mtg/v2/mtglib/internal/tls"
|
||||
"github.com/stretchr/testify/mock"
|
||||
"github.com/stretchr/testify/suite"
|
||||
)
|
||||
|
||||
type ConnMock struct {
|
||||
testlib.EssentialsConnMock
|
||||
|
||||
mu sync.Mutex
|
||||
writeBuffer bytes.Buffer
|
||||
}
|
||||
|
||||
func (m *ConnMock) Write(p []byte) (int, error) {
|
||||
args := m.Called(p)
|
||||
if err := args.Error(1); err != nil {
|
||||
return args.Int(0), err
|
||||
}
|
||||
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
|
||||
return m.writeBuffer.Write(p)
|
||||
}
|
||||
|
||||
func (m *ConnMock) Written() []byte {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
|
||||
return bytes.Clone(m.writeBuffer.Bytes())
|
||||
}
|
||||
|
||||
type ConnTestSuite struct {
|
||||
suite.Suite
|
||||
|
||||
connMock *ConnMock
|
||||
ctx context.Context
|
||||
ctxCancel context.CancelFunc
|
||||
}
|
||||
|
||||
func (suite *ConnTestSuite) SetupTest() {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
suite.ctx = ctx
|
||||
suite.ctxCancel = cancel
|
||||
suite.connMock = &ConnMock{}
|
||||
}
|
||||
|
||||
func (suite *ConnTestSuite) TearDownTest() {
|
||||
suite.ctxCancel()
|
||||
suite.connMock.AssertExpectations(suite.T())
|
||||
}
|
||||
|
||||
func (suite *ConnTestSuite) makeConn() Conn {
|
||||
return NewConn(suite.ctx, suite.connMock, &Stats{
|
||||
k: 2.0,
|
||||
lambda: 0.01,
|
||||
})
|
||||
}
|
||||
|
||||
func (suite *ConnTestSuite) TestWriteBuffersData() {
|
||||
suite.connMock.
|
||||
On("Write", mock.AnythingOfType("[]uint8")).
|
||||
Return(0, nil).
|
||||
Maybe()
|
||||
|
||||
c := suite.makeConn()
|
||||
defer c.Stop()
|
||||
|
||||
n, err := c.Write([]byte{1, 2, 3})
|
||||
suite.NoError(err)
|
||||
suite.Equal(3, n)
|
||||
}
|
||||
|
||||
func (suite *ConnTestSuite) TestWriteOutputsTLSRecords() {
|
||||
suite.connMock.
|
||||
On("Write", mock.AnythingOfType("[]uint8")).
|
||||
Return(0, nil).
|
||||
Maybe()
|
||||
|
||||
c := suite.makeConn()
|
||||
|
||||
payload := []byte("hello doppelganger")
|
||||
_, err := c.Write(payload)
|
||||
suite.NoError(err)
|
||||
|
||||
suite.Eventually(func() bool {
|
||||
return len(suite.connMock.Written()) > 0
|
||||
}, 2*time.Second, time.Millisecond)
|
||||
|
||||
c.Stop()
|
||||
|
||||
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])
|
||||
suite.Equal(tls.TLSVersion[:], header[tls.SizeRecordType:tls.SizeRecordType+tls.SizeVersion])
|
||||
|
||||
length := binary.BigEndian.Uint16(header[tls.SizeRecordType+tls.SizeVersion:])
|
||||
suite.Greater(length, uint16(0))
|
||||
|
||||
rec := make([]byte, length)
|
||||
_, err := io.ReadFull(reader, rec)
|
||||
suite.NoError(err)
|
||||
|
||||
assembled.Write(rec)
|
||||
}
|
||||
|
||||
suite.Equal(payload, assembled.Bytes())
|
||||
}
|
||||
|
||||
func (suite *ConnTestSuite) TestWriteReturnsErrorAfterStop() {
|
||||
suite.connMock.
|
||||
On("Write", mock.AnythingOfType("[]uint8")).
|
||||
Return(0, nil).
|
||||
Maybe()
|
||||
|
||||
c := suite.makeConn()
|
||||
c.Stop()
|
||||
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
|
||||
_, err := c.Write([]byte{1})
|
||||
suite.Error(err)
|
||||
}
|
||||
|
||||
func (suite *ConnTestSuite) TestStopOnUnderlyingWriteError() {
|
||||
suite.connMock.
|
||||
On("Write", mock.AnythingOfType("[]uint8")).
|
||||
Return(0, errors.New("connection reset")).
|
||||
Maybe()
|
||||
|
||||
c := suite.makeConn()
|
||||
|
||||
_, _ = c.Write([]byte("data"))
|
||||
|
||||
suite.Eventually(func() bool {
|
||||
_, err := c.Write([]byte{1})
|
||||
return err != nil
|
||||
}, 2*time.Second, time.Millisecond)
|
||||
}
|
||||
|
||||
func TestConn(t *testing.T) {
|
||||
t.Parallel()
|
||||
suite.Run(t, &ConnTestSuite{})
|
||||
}
|
||||
@@ -0,0 +1,173 @@
|
||||
package doppel
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/9seconds/mtg/v2/essentials"
|
||||
)
|
||||
|
||||
const (
|
||||
DoppelGangerMaxDurations = 4096
|
||||
DoppelGangerScoutMissionEach = 30 * time.Minute
|
||||
DoppelGangerScoutRepeats = 10
|
||||
)
|
||||
|
||||
type gangerConnRequest struct {
|
||||
ret chan Conn
|
||||
payload essentials.Conn
|
||||
}
|
||||
|
||||
type Ganger struct {
|
||||
ctx context.Context
|
||||
ctxCancel context.CancelFunc
|
||||
logger Logger
|
||||
wg sync.WaitGroup
|
||||
|
||||
scout Scout
|
||||
scoutMissionEach time.Duration
|
||||
scoutMissionRepeats int
|
||||
|
||||
stats *Stats
|
||||
durations []time.Duration
|
||||
|
||||
connRequests chan gangerConnRequest
|
||||
}
|
||||
|
||||
func (g *Ganger) Shutdown() {
|
||||
g.ctxCancel()
|
||||
g.wg.Wait()
|
||||
}
|
||||
|
||||
func (g *Ganger) Run() {
|
||||
g.wg.Go(func() {
|
||||
g.run()
|
||||
})
|
||||
}
|
||||
|
||||
func (g *Ganger) NewConn(conn essentials.Conn) (Conn, error) {
|
||||
req := gangerConnRequest{
|
||||
ret: make(chan Conn),
|
||||
payload: conn,
|
||||
}
|
||||
defer close(req.ret)
|
||||
|
||||
select {
|
||||
case <-g.ctx.Done():
|
||||
return Conn{}, context.Cause(g.ctx)
|
||||
case g.connRequests <- req:
|
||||
}
|
||||
|
||||
select {
|
||||
case <-g.ctx.Done():
|
||||
return Conn{}, context.Cause(g.ctx)
|
||||
case conn := <-req.ret:
|
||||
return conn, nil
|
||||
}
|
||||
}
|
||||
|
||||
func (g *Ganger) run() {
|
||||
scoutTicker := time.NewTicker(g.scoutMissionEach)
|
||||
defer func() {
|
||||
scoutTicker.Stop()
|
||||
|
||||
select {
|
||||
case <-scoutTicker.C:
|
||||
default:
|
||||
}
|
||||
}()
|
||||
|
||||
scoutCollectedChan := make(chan []time.Duration)
|
||||
currentScoutCollectedChan := scoutCollectedChan
|
||||
|
||||
updatedStatsChan := make(chan *Stats)
|
||||
|
||||
g.wg.Go(func() {
|
||||
g.runScoutMission(scoutCollectedChan)
|
||||
})
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-g.ctx.Done():
|
||||
return
|
||||
case durations := <-currentScoutCollectedChan:
|
||||
g.durations = append(g.durations, durations...)
|
||||
if len(g.durations) > DoppelGangerMaxDurations {
|
||||
g.durations = g.durations[len(g.durations)-DoppelGangerMaxDurations:]
|
||||
}
|
||||
|
||||
currentScoutCollectedChan = nil
|
||||
g.wg.Go(func() {
|
||||
select {
|
||||
case <-g.ctx.Done():
|
||||
case updatedStatsChan <- NewStats(durations):
|
||||
}
|
||||
})
|
||||
case stats := <-updatedStatsChan:
|
||||
g.stats = stats
|
||||
currentScoutCollectedChan = scoutCollectedChan
|
||||
case <-scoutTicker.C:
|
||||
g.wg.Go(func() {
|
||||
g.runScoutMission(scoutCollectedChan)
|
||||
})
|
||||
case req := <-g.connRequests:
|
||||
select {
|
||||
case <-g.ctx.Done():
|
||||
case req.ret <- NewConn(g.ctx, req.payload, g.stats):
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (g *Ganger) runScoutMission(rvChan chan<- []time.Duration) {
|
||||
durations := []time.Duration{}
|
||||
|
||||
for range g.scoutMissionRepeats {
|
||||
learned, err := g.scout.Learn(g.ctx)
|
||||
if err != nil {
|
||||
g.logger.WarningError("cannot learn", err)
|
||||
continue
|
||||
}
|
||||
durations = append(durations, learned...)
|
||||
}
|
||||
|
||||
select {
|
||||
case <-g.ctx.Done():
|
||||
return
|
||||
case rvChan <- durations:
|
||||
}
|
||||
}
|
||||
|
||||
func NewGanger(
|
||||
ctx context.Context,
|
||||
network Network,
|
||||
logger Logger,
|
||||
scoutEach time.Duration,
|
||||
scoutRepeats int,
|
||||
urls []string,
|
||||
) *Ganger {
|
||||
ctx, cancel := context.WithCancel(ctx)
|
||||
|
||||
if scoutEach == 0 {
|
||||
scoutEach = DoppelGangerScoutMissionEach
|
||||
}
|
||||
|
||||
if scoutRepeats == 0 {
|
||||
scoutRepeats = DoppelGangerScoutRepeats
|
||||
}
|
||||
|
||||
return &Ganger{
|
||||
ctx: ctx,
|
||||
ctxCancel: cancel,
|
||||
logger: logger,
|
||||
scoutMissionEach: scoutEach,
|
||||
scoutMissionRepeats: scoutRepeats,
|
||||
stats: &Stats{
|
||||
k: StatsDefaultK,
|
||||
lambda: StatsDefaultLambda,
|
||||
},
|
||||
scout: NewScout(network, urls),
|
||||
connRequests: make(chan gangerConnRequest),
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,107 @@
|
||||
package doppel
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/9seconds/mtg/v2/internal/testlib"
|
||||
"github.com/stretchr/testify/mock"
|
||||
"github.com/stretchr/testify/suite"
|
||||
)
|
||||
|
||||
type GangerTestSuite struct {
|
||||
TLSServerTestSuite
|
||||
|
||||
log *LoggerMock
|
||||
g *Ganger
|
||||
}
|
||||
|
||||
func (suite *GangerTestSuite) SetupTest() {
|
||||
suite.TLSServerTestSuite.SetupTest()
|
||||
|
||||
suite.log = &LoggerMock{}
|
||||
suite.log.
|
||||
On("Info", mock.AnythingOfType("string")).
|
||||
Maybe()
|
||||
suite.log.
|
||||
On("WarningError", mock.AnythingOfType("string"), mock.Anything).
|
||||
Maybe()
|
||||
|
||||
suite.g = NewGanger(suite.ctx, suite.network, suite.log, time.Hour, 1, suite.urls)
|
||||
suite.g.Run()
|
||||
}
|
||||
|
||||
func (suite *GangerTestSuite) TearDownTest() {
|
||||
suite.g.Shutdown()
|
||||
|
||||
suite.log.AssertExpectations(suite.T())
|
||||
suite.TLSServerTestSuite.TearDownTest()
|
||||
}
|
||||
|
||||
func (suite *GangerTestSuite) TestNewConnAfterShutdown() {
|
||||
suite.g.Shutdown()
|
||||
connMock := &testlib.EssentialsConnMock{}
|
||||
|
||||
_, err := suite.g.NewConn(connMock)
|
||||
suite.Error(err)
|
||||
}
|
||||
|
||||
func (suite *GangerTestSuite) TestNewConnWhileRunning() {
|
||||
connMock := &testlib.EssentialsConnMock{}
|
||||
connMock.
|
||||
On("Write", mock.AnythingOfType("[]uint8")).
|
||||
Return(0, nil).
|
||||
Maybe()
|
||||
connMock.On("Close").
|
||||
Return(nil).
|
||||
Maybe()
|
||||
|
||||
conn, err := suite.g.NewConn(connMock)
|
||||
suite.NoError(err)
|
||||
|
||||
conn.Stop()
|
||||
}
|
||||
|
||||
func (suite *GangerTestSuite) TestNewConnWriteProducesTLSRecords() {
|
||||
var (
|
||||
mu sync.Mutex
|
||||
buf bytes.Buffer
|
||||
)
|
||||
|
||||
connMock := &testlib.EssentialsConnMock{}
|
||||
connMock.On("Write", mock.AnythingOfType("[]uint8")).
|
||||
Run(func(args mock.Arguments) {
|
||||
mu.Lock()
|
||||
buf.Write(args.Get(0).([]byte))
|
||||
mu.Unlock()
|
||||
}).
|
||||
Return(0, nil).
|
||||
Maybe()
|
||||
connMock.On("Close").
|
||||
Return(nil).
|
||||
Maybe()
|
||||
|
||||
conn, err := suite.g.NewConn(connMock)
|
||||
suite.NoError(err)
|
||||
|
||||
payload := bytes.Repeat([]byte("x"), 512)
|
||||
_, err = conn.Write(payload)
|
||||
suite.NoError(err)
|
||||
|
||||
time.Sleep(500 * time.Millisecond)
|
||||
conn.Stop()
|
||||
|
||||
mu.Lock()
|
||||
written := buf.Bytes()
|
||||
mu.Unlock()
|
||||
|
||||
suite.NotEmpty(written)
|
||||
}
|
||||
|
||||
func TestGanger(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
suite.Run(t, &GangerTestSuite{})
|
||||
}
|
||||
@@ -0,0 +1,38 @@
|
||||
package doppel
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"github.com/9seconds/mtg/v2/essentials"
|
||||
"github.com/9seconds/mtg/v2/mtglib/internal/tls"
|
||||
)
|
||||
|
||||
const (
|
||||
// Please see Stats description
|
||||
// https://blog.cloudflare.com/optimizing-tls-over-tcp-to-reduce-latency/
|
||||
// https://github.com/cloudflare/sslconfig/blob/master/patches/nginx__dynamic_tls_records.patch
|
||||
TLSRecordSizeStart = 1369
|
||||
TLSRecordSizeAccel = 4229
|
||||
TLSRecordSizeMax = 16384 - tls.SizeHeader
|
||||
|
||||
TLSCounterAccelAfter = 40
|
||||
TLSCounterMaxAfter = TLSCounterAccelAfter + 20
|
||||
|
||||
TLSRecordSizeResetAfter = time.Second
|
||||
)
|
||||
|
||||
// copypasted from mtglib
|
||||
type Network interface {
|
||||
// Dial establishes context-free TCP connections.
|
||||
Dial(network, address string) (essentials.Conn, error)
|
||||
|
||||
// DialContext dials using a context. This is a preferrable way of
|
||||
// establishing TCP connections.
|
||||
DialContext(ctx context.Context, network, address string) (essentials.Conn, error)
|
||||
|
||||
// MakeHTTPClient build an HTTP client with given dial function. If nothing is
|
||||
// provided, then DialContext of this interface is going to be used.
|
||||
MakeHTTPClient(func(ctx context.Context, network, address string) (essentials.Conn, error)) *http.Client
|
||||
}
|
||||
@@ -0,0 +1,104 @@
|
||||
package doppel
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/tls"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"time"
|
||||
|
||||
"github.com/9seconds/mtg/v2/essentials"
|
||||
"github.com/stretchr/testify/mock"
|
||||
"github.com/stretchr/testify/suite"
|
||||
)
|
||||
|
||||
type SimpleNetwork struct {
|
||||
}
|
||||
|
||||
func (s SimpleNetwork) Dial(network, address string) (essentials.Conn, error) {
|
||||
return s.DialContext(context.Background(), network, address)
|
||||
}
|
||||
|
||||
func (s SimpleNetwork) DialContext(ctx context.Context, network, address string) (essentials.Conn, error) {
|
||||
d := &net.Dialer{}
|
||||
|
||||
conn, err := d.DialContext(ctx, network, address)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return conn.(*net.TCPConn), nil
|
||||
}
|
||||
|
||||
func (s SimpleNetwork) MakeHTTPClient(dialFunc func(ctx context.Context, network, address string) (essentials.Conn, error)) *http.Client {
|
||||
if dialFunc == nil {
|
||||
dialFunc = s.DialContext
|
||||
}
|
||||
|
||||
return &http.Client{
|
||||
Transport: &http.Transport{
|
||||
TLSClientConfig: &tls.Config{
|
||||
InsecureSkipVerify: true, //nolint: gosec
|
||||
},
|
||||
DialContext: func(ctx context.Context, network, address string) (net.Conn, error) {
|
||||
return dialFunc(ctx, network, address)
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
type TLSServerTestSuite struct {
|
||||
suite.Suite
|
||||
|
||||
tlsServer *httptest.Server
|
||||
ctx context.Context
|
||||
ctxCancel context.CancelFunc
|
||||
network SimpleNetwork
|
||||
urls []string
|
||||
}
|
||||
|
||||
func (suite *TLSServerTestSuite) SetupSuite() {
|
||||
suite.tlsServer = httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
|
||||
w.WriteHeader(http.StatusOK)
|
||||
w.Header().Add("Hello", "how long")
|
||||
|
||||
if _, err := w.Write([]byte{1, 2, 3}); err != nil {
|
||||
panic(err)
|
||||
}
|
||||
|
||||
time.Sleep(5 * time.Millisecond)
|
||||
|
||||
if _, err := w.Write([]byte{1, 2, 3}); err != nil {
|
||||
panic(err)
|
||||
}
|
||||
}))
|
||||
suite.urls = []string{suite.tlsServer.URL}
|
||||
}
|
||||
|
||||
func (suite *TLSServerTestSuite) SetupTest() {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
suite.ctx = ctx
|
||||
suite.ctxCancel = cancel
|
||||
}
|
||||
|
||||
func (suite *TLSServerTestSuite) TearDownTest() {
|
||||
suite.ctxCancel()
|
||||
suite.tlsServer.CloseClientConnections()
|
||||
}
|
||||
|
||||
func (suite *TLSServerTestSuite) TearDownSuite() {
|
||||
suite.tlsServer.Close()
|
||||
}
|
||||
|
||||
type LoggerMock struct {
|
||||
mock.Mock
|
||||
}
|
||||
|
||||
func (l *LoggerMock) Info(msg string) {
|
||||
l.Called(msg)
|
||||
}
|
||||
|
||||
func (l *LoggerMock) WarningError(msg string, err error) {
|
||||
l.Called(msg, err)
|
||||
}
|
||||
@@ -0,0 +1,6 @@
|
||||
package doppel
|
||||
|
||||
type Logger interface {
|
||||
Info(msg string)
|
||||
WarningError(msg string, err error)
|
||||
}
|
||||
@@ -0,0 +1,104 @@
|
||||
package doppel
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/9seconds/mtg/v2/essentials"
|
||||
"github.com/9seconds/mtg/v2/mtglib/internal/tls"
|
||||
)
|
||||
|
||||
type Scout struct {
|
||||
network Network
|
||||
urls []string
|
||||
}
|
||||
|
||||
func (s Scout) Learn(ctx context.Context) ([]time.Duration, error) {
|
||||
var durations []time.Duration
|
||||
|
||||
for _, url := range s.urls {
|
||||
learned, err := s.learn(ctx, url)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
durations = append(durations, learned...)
|
||||
}
|
||||
|
||||
return durations, nil
|
||||
}
|
||||
|
||||
func (s Scout) learn(ctx context.Context, url string) ([]time.Duration, error) {
|
||||
client, results := s.makeClient()
|
||||
|
||||
if !strings.HasPrefix(url, "https://") {
|
||||
return nil, fmt.Errorf("url %s must be https", url)
|
||||
}
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
resp, err := client.Do(req)
|
||||
if resp != nil {
|
||||
io.Copy(io.Discard, resp.Body) //nolint: errcheck
|
||||
resp.Body.Close() //nolint: errcheck
|
||||
client.CloseIdleConnections()
|
||||
}
|
||||
|
||||
if err != nil || len(results.data) == 0 {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
durations := []time.Duration{}
|
||||
lastTimestamp := time.Time{}
|
||||
|
||||
for i, v := range results.data {
|
||||
if v.recordType != tls.TypeApplicationData {
|
||||
continue
|
||||
}
|
||||
|
||||
if lastTimestamp.IsZero() {
|
||||
if i > 0 {
|
||||
lastTimestamp = results.data[i-1].timestamp
|
||||
} else {
|
||||
lastTimestamp = v.timestamp
|
||||
}
|
||||
}
|
||||
|
||||
durations = append(durations, v.timestamp.Sub(lastTimestamp))
|
||||
lastTimestamp = v.timestamp
|
||||
}
|
||||
|
||||
return durations, nil
|
||||
}
|
||||
|
||||
func (s Scout) makeClient() (*http.Client, *ScoutConnCollected) {
|
||||
collected := NewScoutConnCollected()
|
||||
client := s.network.MakeHTTPClient(func(
|
||||
ctx context.Context,
|
||||
network string,
|
||||
address string,
|
||||
) (essentials.Conn, error) {
|
||||
conn, err := s.network.DialContext(ctx, network, address)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return NewScoutConn(conn, collected), nil
|
||||
})
|
||||
|
||||
return client, collected
|
||||
}
|
||||
|
||||
func NewScout(network Network, urls []string) Scout {
|
||||
return Scout{
|
||||
network: network,
|
||||
urls: urls,
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,57 @@
|
||||
package doppel
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/binary"
|
||||
"io"
|
||||
|
||||
"github.com/9seconds/mtg/v2/essentials"
|
||||
"github.com/9seconds/mtg/v2/mtglib/internal/tls"
|
||||
)
|
||||
|
||||
type ScoutConn struct {
|
||||
tls.Conn
|
||||
|
||||
results *ScoutConnCollected
|
||||
rawBuf *bytes.Buffer
|
||||
}
|
||||
|
||||
func (s ScoutConn) Read(p []byte) (int, error) {
|
||||
buf := &bytes.Buffer{}
|
||||
|
||||
for {
|
||||
if n, err := s.rawBuf.Read(p); err == nil {
|
||||
return n, nil
|
||||
}
|
||||
|
||||
s.rawBuf.Reset()
|
||||
|
||||
recordType, length, err := tls.ReadRecord(s.Conn, buf)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
|
||||
s.results.Add(recordType)
|
||||
s.rawBuf.Write([]byte{recordType})
|
||||
s.rawBuf.Write(tls.TLSVersion[:])
|
||||
|
||||
if err := binary.Write(s.rawBuf, binary.BigEndian, uint16(length)); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
|
||||
if _, err := io.Copy(s.rawBuf, buf); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func NewScoutConn(conn essentials.Conn, results *ScoutConnCollected) ScoutConn {
|
||||
rawBuf := &bytes.Buffer{}
|
||||
rawBuf.Grow(tls.MaxRecordSize)
|
||||
|
||||
return ScoutConn{
|
||||
Conn: tls.New(conn, false, false),
|
||||
results: results,
|
||||
rawBuf: rawBuf,
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,29 @@
|
||||
package doppel
|
||||
|
||||
import "time"
|
||||
|
||||
const (
|
||||
ScoutConnCollectedPreallocSize = 100
|
||||
)
|
||||
|
||||
type ScoutConnResult struct {
|
||||
timestamp time.Time
|
||||
recordType byte
|
||||
}
|
||||
|
||||
type ScoutConnCollected struct {
|
||||
data []ScoutConnResult
|
||||
}
|
||||
|
||||
func (s *ScoutConnCollected) Add(record byte) {
|
||||
s.data = append(s.data, ScoutConnResult{
|
||||
timestamp: time.Now(),
|
||||
recordType: record,
|
||||
})
|
||||
}
|
||||
|
||||
func NewScoutConnCollected() *ScoutConnCollected {
|
||||
return &ScoutConnCollected{
|
||||
data: make([]ScoutConnResult, 0, ScoutConnCollectedPreallocSize),
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,42 @@
|
||||
package doppel
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/9seconds/mtg/v2/mtglib/internal/tls"
|
||||
"github.com/stretchr/testify/suite"
|
||||
)
|
||||
|
||||
type ScoutConnCollectedTestSuite struct {
|
||||
suite.Suite
|
||||
}
|
||||
|
||||
func (suite *ScoutConnCollectedTestSuite) TestAddSingle() {
|
||||
collected := NewScoutConnCollected()
|
||||
collected.Add(tls.TypeApplicationData)
|
||||
|
||||
suite.Len(collected.data, 1)
|
||||
suite.Equal(byte(tls.TypeApplicationData), collected.data[0].recordType)
|
||||
}
|
||||
|
||||
func (suite *ScoutConnCollectedTestSuite) TestAddTimestampsAreMonotonic() {
|
||||
collected := NewScoutConnCollected()
|
||||
|
||||
collected.Add(tls.TypeApplicationData)
|
||||
|
||||
time.Sleep(time.Microsecond)
|
||||
collected.Add(tls.TypeApplicationData)
|
||||
|
||||
time.Sleep(time.Microsecond)
|
||||
collected.Add(tls.TypeApplicationData)
|
||||
|
||||
for i := 1; i < len(collected.data); i++ {
|
||||
suite.True(collected.data[i].timestamp.After(collected.data[i-1].timestamp))
|
||||
}
|
||||
}
|
||||
|
||||
func TestScoutConnCollected(t *testing.T) {
|
||||
t.Parallel()
|
||||
suite.Run(t, &ScoutConnCollectedTestSuite{})
|
||||
}
|
||||
@@ -0,0 +1,39 @@
|
||||
package doppel
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/suite"
|
||||
)
|
||||
|
||||
type ScoutTestSuite struct {
|
||||
TLSServerTestSuite
|
||||
|
||||
scout Scout
|
||||
}
|
||||
|
||||
func (suite *ScoutTestSuite) SetupSuite() {
|
||||
suite.TLSServerTestSuite.SetupSuite()
|
||||
|
||||
suite.scout = Scout{
|
||||
network: suite.network,
|
||||
urls: suite.urls,
|
||||
}
|
||||
}
|
||||
|
||||
func (suite *ScoutTestSuite) TestCollectResults() {
|
||||
durations, err := suite.scout.Learn(suite.ctx)
|
||||
suite.NoError(err)
|
||||
suite.Less(3, len(durations))
|
||||
}
|
||||
|
||||
func (suite *ScoutTestSuite) TestCollectNothing() {
|
||||
suite.ctxCancel()
|
||||
|
||||
_, err := suite.scout.Learn(suite.ctx)
|
||||
suite.Error(err)
|
||||
}
|
||||
|
||||
func TestScout(t *testing.T) {
|
||||
suite.Run(t, &ScoutTestSuite{})
|
||||
}
|
||||
@@ -0,0 +1,150 @@
|
||||
package doppel
|
||||
|
||||
import (
|
||||
"math"
|
||||
"math/rand/v2"
|
||||
"time"
|
||||
)
|
||||
|
||||
const (
|
||||
StatsBisectTimes = 70
|
||||
StatsLowK = 0.01
|
||||
StatsHighK = 10.0
|
||||
|
||||
StatsDefaultK = 0.6
|
||||
StatsDefaultLambda = 0.002
|
||||
)
|
||||
|
||||
// Stats is responsible for generating values that are distributed according
|
||||
// to some statistical distribution.
|
||||
//
|
||||
// It follows several ideas:
|
||||
// 1. Based on nginx and Cloudflare behaviour, even if server is eager
|
||||
// to send a lot, they all start with small TLS packets that are
|
||||
// approximately MTU-sized. After
|
||||
// 2. After ~40 TLS records, server considers TCP session as somewhat solid
|
||||
// and reliable and ramps up to 4096.
|
||||
// 3. After ~20 TLS records more it jumps to the max 16384 bytes and keep
|
||||
// this size as long as it can
|
||||
// 4. If there is no any byte within a connection for a longer time period,
|
||||
// this counter resets.
|
||||
//
|
||||
// This is called Dynamic TLS Record Sizing
|
||||
// - https://blog.cloudflare.com/optimizing-tls-over-tcp-to-reduce-latency/
|
||||
// - https://community.f5.com/kb/technicalarticles/boosting-tls-performance-with-dynamic-record-sizing-on-big-ip/280798
|
||||
// - https://www.igvita.com/2013/10/24/optimizing-tls-record-size-and-buffering-latency/
|
||||
//
|
||||
// And this optimized for the very first byte, so web browsers could start to
|
||||
// render as early as possible, showing user some preliminary results, optimizing
|
||||
// for perceived latency.
|
||||
//
|
||||
// Since this is very typical for the website, we also aim for that.
|
||||
//
|
||||
// Another important idea is how delays between TLS packets are distributed.
|
||||
// In case of sending huge heavy content with max sized record, delays have
|
||||
// lognormal distribution. But a nature of a typical website shows that
|
||||
// it eagers to deliver as fast as it can in a few very first records and
|
||||
// could possibly slow down later.
|
||||
//
|
||||
// This is perfectly described by Weibull distribution:
|
||||
// - https://en.wikipedia.org/wiki/Weibull_distribution
|
||||
// - https://ieeexplore.ieee.org/document/6662948
|
||||
// - https://www.researchgate.net/publication/224621285_Traffic_modelling_and_cost_optimization_for_transmitting_traffic_messages_over_a_hybrid_broadcast_and_cellular_network
|
||||
// - https://ir.uitm.edu.my/id/eprint/105386/1/105386.pdf
|
||||
//
|
||||
// In other word, a combination of Dynamic TLS Record Sizing hints us for
|
||||
// Weibull distribution.
|
||||
type Stats struct {
|
||||
sizeLastRequested time.Time
|
||||
sizeCounter int
|
||||
|
||||
// https://en.wikipedia.org/wiki/Shape_parameter
|
||||
k float64
|
||||
// https://en.wikipedia.org/wiki/Scale_parameter
|
||||
lambda float64
|
||||
}
|
||||
|
||||
func (d *Stats) Delay() time.Duration {
|
||||
// u ∈ (0, 1], avoids ln(0)
|
||||
u := 1.0 - rand.Float64()
|
||||
|
||||
// X = λ·(-ln U)^(1/k)
|
||||
generated := d.lambda * math.Pow(-math.Log(u), 1.0/d.k)
|
||||
|
||||
// generated is in milliseconds
|
||||
return time.Duration(generated * float64(time.Millisecond))
|
||||
}
|
||||
|
||||
func (d *Stats) Size() int {
|
||||
if time.Since(d.sizeLastRequested) > TLSRecordSizeResetAfter {
|
||||
d.sizeCounter = 0
|
||||
}
|
||||
|
||||
d.sizeLastRequested = time.Now()
|
||||
d.sizeCounter++
|
||||
|
||||
switch {
|
||||
case d.sizeCounter <= TLSCounterAccelAfter:
|
||||
return TLSRecordSizeStart
|
||||
case d.sizeCounter <= TLSCounterMaxAfter:
|
||||
return TLSRecordSizeAccel
|
||||
}
|
||||
|
||||
return TLSRecordSizeMax
|
||||
}
|
||||
|
||||
func NewStats(durations []time.Duration) *Stats {
|
||||
n := float64(len(durations))
|
||||
|
||||
// in milliseconds
|
||||
durFloats := make([]float64, len(durations))
|
||||
for i, v := range durations {
|
||||
durFloats[i] = float64(v.Microseconds()) / 1000.0
|
||||
}
|
||||
|
||||
// The bisection solves the standard Weibull MLE equation for shape
|
||||
// parameter k. There is no any good formula for doing that so we
|
||||
// approximate it by several bisections. The number of operations
|
||||
// is statically defined by a constant.
|
||||
|
||||
sumLog := 0.0
|
||||
for _, v := range durFloats {
|
||||
sumLog += math.Log(v)
|
||||
}
|
||||
|
||||
lowK := StatsLowK
|
||||
highK := StatsHighK
|
||||
|
||||
for range StatsBisectTimes {
|
||||
midK := (lowK + highK) / 2.0
|
||||
sumXK := 0.0
|
||||
sumXKLog := 0.0
|
||||
|
||||
for _, v := range durFloats {
|
||||
xk := math.Pow(v, midK)
|
||||
sumXK += xk
|
||||
sumXKLog += xk * math.Log(v)
|
||||
}
|
||||
|
||||
if (1.0/midK)+(sumLog/n)-(sumXKLog/sumXK) > 0 {
|
||||
lowK = midK
|
||||
} else {
|
||||
highK = midK
|
||||
}
|
||||
}
|
||||
|
||||
k := (lowK + highK) / 2
|
||||
|
||||
sumXK := 0.0
|
||||
for _, v := range durFloats {
|
||||
sumXK += math.Pow(v, k)
|
||||
}
|
||||
|
||||
// λ = (Σxᵢᵏ / n)^(1/k)
|
||||
lambda := math.Pow(sumXK/n, 1.0/k)
|
||||
|
||||
return &Stats{
|
||||
k: k,
|
||||
lambda: lambda,
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,194 @@
|
||||
package doppel
|
||||
|
||||
import (
|
||||
"math"
|
||||
"math/rand/v2"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/suite"
|
||||
)
|
||||
|
||||
type StatsTestSuite struct {
|
||||
suite.Suite
|
||||
}
|
||||
|
||||
func (suite *StatsTestSuite) GenWeibull(k, lambda float64, n int, seed uint64) []time.Duration {
|
||||
rng := rand.New(rand.NewPCG(seed, 0))
|
||||
samples := make([]time.Duration, n)
|
||||
|
||||
for i := range samples {
|
||||
u := 1.0 - rng.Float64()
|
||||
ms := lambda * math.Pow(-math.Log(u), 1.0/k)
|
||||
d := time.Duration(ms * float64(time.Millisecond))
|
||||
|
||||
if d < time.Microsecond {
|
||||
time.Sleep(time.Microsecond)
|
||||
d = time.Microsecond
|
||||
}
|
||||
|
||||
samples[i] = d
|
||||
}
|
||||
|
||||
return samples
|
||||
}
|
||||
|
||||
func (suite *StatsTestSuite) TestNewStatsRecoverParameters() {
|
||||
knownK := 1.5
|
||||
knownLambda := 100.0
|
||||
|
||||
samples := suite.GenWeibull(knownK, knownLambda, 5000, 42)
|
||||
stats := NewStats(samples)
|
||||
|
||||
suite.InDelta(knownK, stats.k, 0.1)
|
||||
suite.InDelta(knownLambda, stats.lambda, 5.0)
|
||||
}
|
||||
|
||||
func (suite *StatsTestSuite) TestNewStatsExponentialCase() {
|
||||
// When k=1, Weibull reduces to exponential distribution.
|
||||
knownK := 1.0
|
||||
knownLambda := 50.0
|
||||
|
||||
samples := suite.GenWeibull(knownK, knownLambda, 5000, 123)
|
||||
stats := NewStats(samples)
|
||||
|
||||
suite.InDelta(knownK, stats.k, 0.1)
|
||||
suite.InDelta(knownLambda, stats.lambda, 5.0)
|
||||
}
|
||||
|
||||
func (suite *StatsTestSuite) TestNewStatsSmallK() {
|
||||
// k < 1 produces a heavy-tailed distribution typical for network delays.
|
||||
// Lambda must be large enough so samples stay above microsecond precision
|
||||
// after time.Duration round-trip.
|
||||
knownK := 0.6
|
||||
knownLambda := 100.0
|
||||
|
||||
samples := suite.GenWeibull(knownK, knownLambda, 10000, 99)
|
||||
stats := NewStats(samples)
|
||||
|
||||
suite.InDelta(knownK, stats.k, 0.05)
|
||||
suite.InDelta(knownLambda, stats.lambda, 5.0)
|
||||
}
|
||||
|
||||
func (suite *StatsTestSuite) TestNewStatsLargeK() {
|
||||
// k > 1: light tail, concentrated around the mode.
|
||||
knownK := 5.0
|
||||
knownLambda := 200.0
|
||||
|
||||
samples := suite.GenWeibull(knownK, knownLambda, 5000, 77)
|
||||
stats := NewStats(samples)
|
||||
|
||||
suite.InDelta(knownK, stats.k, 0.3)
|
||||
suite.InDelta(knownLambda, stats.lambda, 5.0)
|
||||
}
|
||||
|
||||
func (suite *StatsTestSuite) TestDelayNonNegative() {
|
||||
stats := &Stats{
|
||||
k: 1.5,
|
||||
lambda: 100.0,
|
||||
}
|
||||
|
||||
for range 200 {
|
||||
dur := stats.Delay()
|
||||
suite.GreaterOrEqual(dur, time.Duration(0))
|
||||
}
|
||||
}
|
||||
|
||||
func (suite *StatsTestSuite) TestDelayDistributionMean() {
|
||||
// Weibull mean = λ · Γ(1 + 1/k)
|
||||
k := 2.0
|
||||
lambda := 50.0
|
||||
stats := &Stats{k: k, lambda: lambda}
|
||||
|
||||
n := 50000
|
||||
sum := 0.0
|
||||
|
||||
for range n {
|
||||
dur := stats.Delay()
|
||||
sum += float64(dur) / float64(time.Millisecond)
|
||||
}
|
||||
|
||||
sampleMean := sum / float64(n)
|
||||
expectedMean := lambda * math.Gamma(1.0+1.0/k)
|
||||
|
||||
suite.InDelta(expectedMean, sampleMean, expectedMean*0.05)
|
||||
}
|
||||
|
||||
func (suite *StatsTestSuite) TestNewStatsRoundTrip() {
|
||||
// Estimate parameters from data, then verify that Delay samples
|
||||
// from the fitted distribution have approximately the same mean.
|
||||
knownK := 1.2
|
||||
knownLambda := 80.0
|
||||
|
||||
samples := suite.GenWeibull(knownK, knownLambda, 5000, 555)
|
||||
stats := NewStats(samples)
|
||||
|
||||
n := 50000
|
||||
sum := 0.0
|
||||
|
||||
for range n {
|
||||
dur := stats.Delay()
|
||||
sum += float64(dur) / float64(time.Millisecond)
|
||||
}
|
||||
|
||||
sampleMean := sum / float64(n)
|
||||
expectedMean := knownLambda * math.Gamma(1.0+1.0/knownK)
|
||||
|
||||
suite.InDelta(expectedMean, sampleMean, expectedMean*0.05)
|
||||
}
|
||||
|
||||
func (suite *StatsTestSuite) TestSizeStartPhase() {
|
||||
stats := &Stats{k: 1.0, lambda: 1.0}
|
||||
|
||||
for range TLSCounterAccelAfter {
|
||||
size := stats.Size()
|
||||
suite.Equal(TLSRecordSizeStart, size)
|
||||
}
|
||||
}
|
||||
|
||||
func (suite *StatsTestSuite) TestSizeAccelPhase() {
|
||||
stats := &Stats{k: 1.0, lambda: 1.0}
|
||||
|
||||
for range TLSCounterAccelAfter {
|
||||
stats.Size()
|
||||
}
|
||||
|
||||
for range TLSCounterMaxAfter - TLSCounterAccelAfter {
|
||||
size := stats.Size()
|
||||
suite.Equal(TLSRecordSizeAccel, size)
|
||||
}
|
||||
}
|
||||
|
||||
func (suite *StatsTestSuite) TestSizeMaxPhase() {
|
||||
stats := &Stats{k: 1.0, lambda: 1.0}
|
||||
|
||||
for range TLSCounterMaxAfter {
|
||||
stats.Size()
|
||||
}
|
||||
|
||||
for range 20 {
|
||||
size := stats.Size()
|
||||
suite.Equal(TLSRecordSizeMax, size)
|
||||
}
|
||||
}
|
||||
|
||||
func (suite *StatsTestSuite) TestSizeResetsAfterInactivity() {
|
||||
stats := &Stats{k: 1.0, lambda: 1.0}
|
||||
|
||||
// Advance past start phase.
|
||||
for range TLSCounterMaxAfter {
|
||||
stats.Size()
|
||||
}
|
||||
|
||||
suite.Equal(TLSRecordSizeMax, stats.Size())
|
||||
|
||||
// Simulate inactivity by backdating sizeLastRequested.
|
||||
stats.sizeLastRequested = time.Now().Add(-TLSRecordSizeResetAfter - time.Millisecond)
|
||||
|
||||
suite.Equal(TLSRecordSizeStart, stats.Size())
|
||||
}
|
||||
|
||||
func TestStats(t *testing.T) {
|
||||
t.Parallel()
|
||||
suite.Run(t, &StatsTestSuite{})
|
||||
}
|
||||
@@ -0,0 +1,88 @@
|
||||
package tls
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
|
||||
"github.com/9seconds/mtg/v2/essentials"
|
||||
)
|
||||
|
||||
const (
|
||||
SizeRecordType = 1
|
||||
SizeVersion = 2
|
||||
SizeSize = 2
|
||||
SizeHeader = SizeRecordType + SizeVersion + SizeSize
|
||||
|
||||
MaxRecordSize = 16384
|
||||
MaxRecordPayloadSize = MaxRecordSize - SizeHeader
|
||||
DefaultBufferSize = 4096
|
||||
|
||||
TypeChangeCipherSpec = 0x14
|
||||
TypeHandshake = 0x16
|
||||
TypeApplicationData = 0x17
|
||||
)
|
||||
|
||||
var (
|
||||
// TLS 1.2 is used for both TLS 1.2 and 1.3
|
||||
TLSVersion = [SizeVersion]byte{3, 3}
|
||||
)
|
||||
|
||||
// Conn presents an established TLS 1.3 connection, after handshake
|
||||
type Conn struct {
|
||||
essentials.Conn
|
||||
|
||||
p *connPayload
|
||||
}
|
||||
|
||||
type connPayload struct {
|
||||
readBuf bytes.Buffer
|
||||
writeBuf bytes.Buffer
|
||||
connBuffered *bufio.Reader
|
||||
read bool
|
||||
write bool
|
||||
}
|
||||
|
||||
func (c Conn) Write(p []byte) (int, error) {
|
||||
if !c.p.write {
|
||||
return c.Conn.Write(p)
|
||||
}
|
||||
|
||||
return len(p), WriteRecord(c.Conn, p)
|
||||
}
|
||||
|
||||
func (c Conn) Read(p []byte) (int, error) {
|
||||
if !c.p.read {
|
||||
return c.Conn.Read(p)
|
||||
}
|
||||
|
||||
for {
|
||||
if n, err := c.p.readBuf.Read(p); err == nil {
|
||||
return n, nil
|
||||
}
|
||||
|
||||
recordType, _, err := ReadRecord(c.p.connBuffered, &c.p.readBuf)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
|
||||
if recordType != TypeApplicationData {
|
||||
c.p.readBuf.Reset()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func New(conn essentials.Conn, read, write bool) Conn {
|
||||
newConn := Conn{
|
||||
Conn: conn,
|
||||
p: &connPayload{
|
||||
connBuffered: bufio.NewReaderSize(conn, DefaultBufferSize),
|
||||
read: read,
|
||||
write: write,
|
||||
},
|
||||
}
|
||||
|
||||
newConn.p.readBuf.Grow(DefaultBufferSize)
|
||||
newConn.p.writeBuf.Grow(DefaultBufferSize)
|
||||
|
||||
return newConn
|
||||
}
|
||||
@@ -0,0 +1,160 @@
|
||||
package tls
|
||||
|
||||
import (
|
||||
"io"
|
||||
"testing"
|
||||
|
||||
"github.com/9seconds/mtg/v2/internal/testlib"
|
||||
"github.com/stretchr/testify/mock"
|
||||
"github.com/stretchr/testify/suite"
|
||||
)
|
||||
|
||||
type ConnTestSuite struct {
|
||||
suite.Suite
|
||||
|
||||
connMock *testlib.EssentialsConnMock
|
||||
}
|
||||
|
||||
func (suite *ConnTestSuite) SetupTest() {
|
||||
suite.connMock = &testlib.EssentialsConnMock{}
|
||||
}
|
||||
|
||||
func (suite *ConnTestSuite) TearDownTest() {
|
||||
suite.connMock.AssertExpectations(suite.T())
|
||||
}
|
||||
|
||||
func (suite *ConnTestSuite) feedRead(raw []byte) {
|
||||
suite.connMock.
|
||||
On("Read", mock.AnythingOfType("[]uint8")).
|
||||
Run(func(args mock.Arguments) {
|
||||
copy(args.Get(0).([]byte), raw)
|
||||
}).
|
||||
Return(len(raw), nil).
|
||||
Once()
|
||||
suite.connMock.
|
||||
On("Read", mock.AnythingOfType("[]uint8")).
|
||||
Return(0, io.EOF).
|
||||
Maybe()
|
||||
}
|
||||
|
||||
func (suite *ConnTestSuite) TestReadTLSEnabled() {
|
||||
payload := []byte("hello world")
|
||||
suite.feedRead(MakeTLSRecord(0x17, payload))
|
||||
|
||||
conn := New(suite.connMock, true, false)
|
||||
|
||||
buf := make([]byte, 128)
|
||||
n, err := conn.Read(buf)
|
||||
|
||||
suite.NoError(err)
|
||||
suite.Equal(payload, buf[:n])
|
||||
}
|
||||
|
||||
func (suite *ConnTestSuite) TestReadTLSSkipsNonApplicationData() {
|
||||
raw := append(
|
||||
MakeTLSRecord(0x14, []byte{1}),
|
||||
MakeTLSRecord(0x17, []byte("real data"))...,
|
||||
)
|
||||
suite.feedRead(raw)
|
||||
|
||||
conn := New(suite.connMock, true, false)
|
||||
|
||||
buf := make([]byte, 128)
|
||||
n, err := conn.Read(buf)
|
||||
|
||||
suite.NoError(err)
|
||||
suite.Equal([]byte("real data"), buf[:n])
|
||||
}
|
||||
|
||||
func (suite *ConnTestSuite) TestReadTLSMultipleRecords() {
|
||||
raw := append(
|
||||
MakeTLSRecord(0x17, []byte("first")),
|
||||
MakeTLSRecord(0x17, []byte("second"))...,
|
||||
)
|
||||
suite.feedRead(raw)
|
||||
|
||||
conn := New(suite.connMock, true, false)
|
||||
buf := make([]byte, 128)
|
||||
|
||||
n, err := conn.Read(buf)
|
||||
suite.NoError(err)
|
||||
suite.Equal([]byte("first"), buf[:n])
|
||||
|
||||
n, err = conn.Read(buf)
|
||||
suite.NoError(err)
|
||||
suite.Equal([]byte("second"), buf[:n])
|
||||
}
|
||||
|
||||
func (suite *ConnTestSuite) TestReadTLSSmallBuffer() {
|
||||
payload := []byte("hello world, this is a longer payload")
|
||||
suite.feedRead(MakeTLSRecord(0x17, payload))
|
||||
|
||||
conn := New(suite.connMock, true, false)
|
||||
|
||||
small := make([]byte, 5)
|
||||
n, err := conn.Read(small)
|
||||
suite.NoError(err)
|
||||
suite.Equal(payload[:5], small[:n])
|
||||
|
||||
rest := make([]byte, 128)
|
||||
n, err = conn.Read(rest)
|
||||
suite.NoError(err)
|
||||
suite.Equal(payload[5:], rest[:n])
|
||||
}
|
||||
|
||||
func (suite *ConnTestSuite) TestReadPassthrough() {
|
||||
data := []byte("raw bytes")
|
||||
|
||||
suite.connMock.
|
||||
On("Read", mock.AnythingOfType("[]uint8")).
|
||||
Run(func(args mock.Arguments) {
|
||||
copy(args.Get(0).([]byte), data)
|
||||
}).
|
||||
Return(len(data), nil).
|
||||
Once()
|
||||
|
||||
conn := New(suite.connMock, false, false)
|
||||
|
||||
buf := make([]byte, 128)
|
||||
n, err := conn.Read(buf)
|
||||
|
||||
suite.NoError(err)
|
||||
suite.Equal(data, buf[:n])
|
||||
}
|
||||
|
||||
func (suite *ConnTestSuite) TestWritePassthrough() {
|
||||
data := []byte("outgoing data")
|
||||
|
||||
suite.connMock.
|
||||
On("Write", mock.AnythingOfType("[]uint8")).
|
||||
Return(len(data), nil).
|
||||
Once()
|
||||
|
||||
conn := New(suite.connMock, false, false)
|
||||
|
||||
n, err := conn.Write(data)
|
||||
|
||||
suite.NoError(err)
|
||||
suite.Equal(len(data), n)
|
||||
}
|
||||
|
||||
func (suite *ConnTestSuite) TestWriteTLSEnabled() {
|
||||
data := []byte("outgoing data")
|
||||
|
||||
suite.connMock.
|
||||
On("Write", mock.AnythingOfType("[]uint8")).
|
||||
Return(len(data), nil).
|
||||
Once()
|
||||
|
||||
conn := New(suite.connMock, false, true)
|
||||
|
||||
n, err := conn.Write(data)
|
||||
|
||||
suite.NoError(err)
|
||||
suite.Equal(len(data), n)
|
||||
}
|
||||
|
||||
func TestConn(t *testing.T) {
|
||||
t.Parallel()
|
||||
suite.Run(t, &ConnTestSuite{})
|
||||
}
|
||||
@@ -0,0 +1,30 @@
|
||||
package tls
|
||||
|
||||
import (
|
||||
"encoding/binary"
|
||||
|
||||
"github.com/stretchr/testify/mock"
|
||||
)
|
||||
|
||||
type WriterMock struct {
|
||||
mock.Mock
|
||||
}
|
||||
|
||||
func (m *WriterMock) Write(p []byte) (int, error) {
|
||||
args := m.Called(p)
|
||||
return args.Int(0), args.Error(1)
|
||||
}
|
||||
|
||||
// makeTLSRecord builds a raw TLS record from hardcoded offsets:
|
||||
// type(1) + version(2, {3,3}) + length(2, big-endian) + payload.
|
||||
func MakeTLSRecord(recordType byte, payload []byte) []byte {
|
||||
buf := make([]byte, 5+len(payload))
|
||||
|
||||
buf[0] = recordType
|
||||
buf[1] = 3
|
||||
buf[2] = 3
|
||||
binary.BigEndian.PutUint16(buf[3:5], uint16(len(payload)))
|
||||
copy(buf[5:], payload)
|
||||
|
||||
return buf
|
||||
}
|
||||
@@ -0,0 +1,48 @@
|
||||
package tls
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/binary"
|
||||
"fmt"
|
||||
"io"
|
||||
)
|
||||
|
||||
func ReadRecord(r io.Reader, w io.Writer) (byte, int64, error) {
|
||||
buf := [SizeHeader]byte{}
|
||||
|
||||
if _, err := io.ReadFull(r, buf[:]); err != nil {
|
||||
return 0, 0, err
|
||||
}
|
||||
|
||||
pVer := buf[SizeRecordType:]
|
||||
pLen := pVer[SizeVersion:]
|
||||
|
||||
if !bytes.Equal(TLSVersion[:], pVer[:SizeVersion]) {
|
||||
return 0, 0, fmt.Errorf("incorrect tls version %v", pVer)
|
||||
}
|
||||
|
||||
length := int64(binary.BigEndian.Uint16(pLen[:SizeSize]))
|
||||
_, err := io.CopyN(w, r, length)
|
||||
|
||||
return buf[0], length, err
|
||||
}
|
||||
|
||||
func WriteRecord(w io.Writer, payload []byte) error {
|
||||
buf := [MaxRecordSize]byte{}
|
||||
buf[0] = TypeApplicationData
|
||||
|
||||
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))
|
||||
}
|
||||
|
||||
_, err := w.Write(buf[:SizeHeader+len(payload)])
|
||||
|
||||
return err
|
||||
}
|
||||
@@ -0,0 +1,125 @@
|
||||
package tls
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/binary"
|
||||
"errors"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/mock"
|
||||
"github.com/stretchr/testify/suite"
|
||||
)
|
||||
|
||||
type UtilsTestSuite struct {
|
||||
suite.Suite
|
||||
|
||||
dst *bytes.Buffer
|
||||
}
|
||||
|
||||
func (suite *UtilsTestSuite) SetupTest() {
|
||||
suite.dst = &bytes.Buffer{}
|
||||
}
|
||||
|
||||
func (suite *UtilsTestSuite) TestReadRecord() {
|
||||
payload := []byte("hello world")
|
||||
raw := MakeTLSRecord(0x17, payload)
|
||||
|
||||
recordType, length, err := ReadRecord(bytes.NewReader(raw), suite.dst)
|
||||
|
||||
suite.NoError(err)
|
||||
suite.Equal(byte(0x17), recordType)
|
||||
suite.Equal(int64(len(payload)), length)
|
||||
suite.Equal(payload, suite.dst.Bytes())
|
||||
}
|
||||
|
||||
func (suite *UtilsTestSuite) TestReadRecordChangeCipherSpec() {
|
||||
payload := []byte{1}
|
||||
raw := MakeTLSRecord(0x14, payload)
|
||||
|
||||
recordType, length, err := ReadRecord(bytes.NewReader(raw), suite.dst)
|
||||
|
||||
suite.NoError(err)
|
||||
suite.Equal(byte(0x14), recordType)
|
||||
suite.Equal(int64(1), length)
|
||||
}
|
||||
|
||||
func (suite *UtilsTestSuite) TestReadRecordRejectsWrongVersion() {
|
||||
record := []byte{0x17, 3, 1, 0, 5, 0, 0, 0, 0, 0}
|
||||
|
||||
_, _, err := ReadRecord(bytes.NewReader(record), suite.dst)
|
||||
suite.ErrorContains(err, "incorrect tls version")
|
||||
}
|
||||
|
||||
func (suite *UtilsTestSuite) TestReadRecordEmptyReader() {
|
||||
_, _, err := ReadRecord(bytes.NewReader(nil), suite.dst)
|
||||
suite.Error(err)
|
||||
}
|
||||
|
||||
func (suite *UtilsTestSuite) TestReadRecordTruncatedHeader() {
|
||||
_, _, err := ReadRecord(bytes.NewReader([]byte{0x17, 3}), suite.dst)
|
||||
suite.Error(err)
|
||||
}
|
||||
|
||||
func (suite *UtilsTestSuite) TestReadRecordTruncatedPayload() {
|
||||
raw := MakeTLSRecord(0x17, []byte("full payload"))
|
||||
truncated := raw[:5+3]
|
||||
|
||||
_, _, err := ReadRecord(bytes.NewReader(truncated), suite.dst)
|
||||
suite.Error(err)
|
||||
}
|
||||
|
||||
func (suite *UtilsTestSuite) TestWriteRecord() {
|
||||
payload := []byte("hello world")
|
||||
|
||||
err := WriteRecord(suite.dst, payload)
|
||||
suite.NoError(err)
|
||||
|
||||
written := suite.dst.Bytes()
|
||||
suite.Equal(byte(0x17), written[0])
|
||||
suite.Equal([]byte{3, 3}, written[1:3])
|
||||
|
||||
length := binary.BigEndian.Uint16(written[3:5])
|
||||
suite.Equal(uint16(len(payload)), length)
|
||||
suite.Equal(payload, written[5:])
|
||||
}
|
||||
|
||||
func (suite *UtilsTestSuite) TestWriteRecordRoundTrip() {
|
||||
payload := []byte("round trip test")
|
||||
|
||||
var wire bytes.Buffer
|
||||
|
||||
err := WriteRecord(&wire, payload)
|
||||
suite.NoError(err)
|
||||
|
||||
var recovered bytes.Buffer
|
||||
|
||||
recordType, length, err := ReadRecord(&wire, &recovered)
|
||||
|
||||
suite.NoError(err)
|
||||
suite.Equal(byte(0x17), recordType)
|
||||
suite.Equal(int64(len(payload)), length)
|
||||
suite.Equal(payload, recovered.Bytes())
|
||||
}
|
||||
|
||||
func (suite *UtilsTestSuite) TestWriteRecordPropagatesError() {
|
||||
m := &WriterMock{}
|
||||
m.
|
||||
On("Write", mock.AnythingOfType("[]uint8")).
|
||||
Once().
|
||||
Return(0, errors.New("dist full"))
|
||||
|
||||
err := WriteRecord(m, []byte("data"))
|
||||
suite.Error(err)
|
||||
|
||||
m.AssertExpectations(suite.T())
|
||||
}
|
||||
|
||||
func (suite *UtilsTestSuite) TestWriteRecordPayloadTooLarge() {
|
||||
err := WriteRecord(suite.dst, make([]byte, MaxRecordPayloadSize+1))
|
||||
suite.Error(err)
|
||||
}
|
||||
|
||||
func TestUtils(t *testing.T) {
|
||||
t.Parallel()
|
||||
suite.Run(t, &UtilsTestSuite{})
|
||||
}
|
||||
Reference in New Issue
Block a user