lib/tlsutil: Allocate UnionedConnection in one go (#7912)

This commit is contained in:
greatroar 2021-09-21 08:40:34 +02:00 committed by GitHub
parent 30e5243f5e
commit a4489dec30
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 19 additions and 18 deletions

View File

@ -185,9 +185,10 @@ func (l *DowngradingListener) AcceptNoWrapTLS() (net.Conn, bool, error) {
return nil, false, err return nil, false, err
} }
var first [1]byte union := &UnionedConnection{Conn: conn}
conn.SetReadDeadline(time.Now().Add(1 * time.Second)) conn.SetReadDeadline(time.Now().Add(1 * time.Second))
n, err := conn.Read(first[:]) n, err := conn.Read(union.first[:])
conn.SetReadDeadline(time.Time{}) conn.SetReadDeadline(time.Time{})
if err != nil || n == 0 { if err != nil || n == 0 {
// We hit a read error here, but the Accept() call succeeded so we must not return an error. // We hit a read error here, but the Accept() call succeeded so we must not return an error.
@ -196,22 +197,23 @@ func (l *DowngradingListener) AcceptNoWrapTLS() (net.Conn, bool, error) {
return conn, false, ErrIdentificationFailed return conn, false, ErrIdentificationFailed
} }
return &UnionedConnection{&first, conn}, first[0] == 0x16, nil return union, union.first[0] == 0x16, nil
} }
type UnionedConnection struct { type UnionedConnection struct {
first *[1]byte first [1]byte
firstDone bool
net.Conn net.Conn
} }
func (c *UnionedConnection) Read(b []byte) (n int, err error) { func (c *UnionedConnection) Read(b []byte) (n int, err error) {
if c.first != nil { if !c.firstDone {
if len(b) == 0 { if len(b) == 0 {
// this probably doesn't happen, but handle it anyway // this probably doesn't happen, but handle it anyway
return 0, nil return 0, nil
} }
b[0] = c.first[0] b[0] = c.first[0]
c.first = nil c.firstDone = true
return 1, nil return 1, nil
} }
return c.Conn.Read(b) return c.Conn.Read(b)

View File

@ -60,9 +60,8 @@ func TestUnionedConnection(t *testing.T) {
if n != 1 { if n != 1 {
t.Errorf("%d: first read returned %d bytes, not 1", i, n) t.Errorf("%d: first read returned %d bytes, not 1", i, n)
} }
// Check that we've nilled out the "first" thing if !conn.(*UnionedConnection).firstDone {
if conn.(*UnionedConnection).first != nil { t.Errorf("%d: expected first read to set the `firstDone` attribute", i)
t.Errorf("%d: expected first read to clear out the `first` attribute", i)
} }
} }
bs = append(bs, buf[:n]...) bs = append(bs, buf[:n]...)
@ -130,8 +129,8 @@ func (f *fakeAccepter) Accept() (net.Conn, error) {
return &fakeConn{f.data}, nil return &fakeConn{f.data}, nil
} }
func (f *fakeAccepter) Addr() net.Addr { return nil } func (*fakeAccepter) Addr() net.Addr { return nil }
func (f *fakeAccepter) Close() error { return nil } func (*fakeAccepter) Close() error { return nil }
type fakeConn struct { type fakeConn struct {
data []byte data []byte
@ -146,13 +145,13 @@ func (f *fakeConn) Read(b []byte) (int, error) {
return n, nil return n, nil
} }
func (f *fakeConn) Write(b []byte) (int, error) { func (*fakeConn) Write(b []byte) (int, error) {
return len(b), nil return len(b), nil
} }
func (f *fakeConn) Close() error { return nil } func (*fakeConn) Close() error { return nil }
func (f *fakeConn) LocalAddr() net.Addr { return nil } func (*fakeConn) LocalAddr() net.Addr { return nil }
func (f *fakeConn) RemoteAddr() net.Addr { return nil } func (*fakeConn) RemoteAddr() net.Addr { return nil }
func (f *fakeConn) SetDeadline(time.Time) error { return nil } func (*fakeConn) SetDeadline(time.Time) error { return nil }
func (f *fakeConn) SetReadDeadline(time.Time) error { return nil } func (*fakeConn) SetReadDeadline(time.Time) error { return nil }
func (f *fakeConn) SetWriteDeadline(time.Time) error { return nil } func (*fakeConn) SetWriteDeadline(time.Time) error { return nil }