mirror of
https://github.com/tailscale/tailscale.git
synced 2024-11-26 03:25:35 +00:00
112 lines
2.6 KiB
Go
112 lines
2.6 KiB
Go
// Copyright (c) 2020 Tailscale Inc & AUTHORS All rights reserved.
|
|
// Use of this source code is governed by a BSD-style
|
|
// license that can be found in the LICENSE file.
|
|
|
|
package nettest
|
|
|
|
import (
|
|
"net"
|
|
"net/netip"
|
|
"time"
|
|
)
|
|
|
|
// Conn is a net.Conn that can additionally have its reads and writes blocked and unblocked.
|
|
type Conn interface {
|
|
net.Conn
|
|
|
|
// SetReadBlock blocks or unblocks the Read method of this Conn.
|
|
// It reports an error if the existing value matches the new value,
|
|
// or if the Conn has been Closed.
|
|
SetReadBlock(bool) error
|
|
|
|
// SetWriteBlock blocks or unblocks the Write method of this Conn.
|
|
// It reports an error if the existing value matches the new value,
|
|
// or if the Conn has been Closed.
|
|
SetWriteBlock(bool) error
|
|
}
|
|
|
|
// NewConn creates a pair of Conns that are wired together by pipes.
|
|
func NewConn(name string, maxBuf int) (Conn, Conn) {
|
|
r := NewPipe(name+"|0", maxBuf)
|
|
w := NewPipe(name+"|1", maxBuf)
|
|
|
|
return &connHalf{r: r, w: w}, &connHalf{r: w, w: r}
|
|
}
|
|
|
|
// NewTCPConn creates a pair of Conns that are wired together by pipes.
|
|
func NewTCPConn(src, dst netip.AddrPort, maxBuf int) (local Conn, remote Conn) {
|
|
r := NewPipe(src.String(), maxBuf)
|
|
w := NewPipe(dst.String(), maxBuf)
|
|
|
|
lAddr := net.TCPAddrFromAddrPort(src)
|
|
rAddr := net.TCPAddrFromAddrPort(dst)
|
|
|
|
return &connHalf{r: r, w: w, remote: rAddr, local: lAddr}, &connHalf{r: w, w: r, remote: lAddr, local: rAddr}
|
|
}
|
|
|
|
type connAddr string
|
|
|
|
func (a connAddr) Network() string { return "mem" }
|
|
func (a connAddr) String() string { return string(a) }
|
|
|
|
type connHalf struct {
|
|
local, remote net.Addr
|
|
r, w *Pipe
|
|
}
|
|
|
|
func (c *connHalf) LocalAddr() net.Addr {
|
|
if c.local != nil {
|
|
return c.local
|
|
}
|
|
return connAddr(c.r.name)
|
|
}
|
|
|
|
func (c *connHalf) RemoteAddr() net.Addr {
|
|
if c.remote != nil {
|
|
return c.remote
|
|
}
|
|
return connAddr(c.w.name)
|
|
}
|
|
|
|
func (c *connHalf) Read(b []byte) (n int, err error) {
|
|
return c.r.Read(b)
|
|
}
|
|
func (c *connHalf) Write(b []byte) (n int, err error) {
|
|
return c.w.Write(b)
|
|
}
|
|
|
|
func (c *connHalf) Close() error {
|
|
if err := c.w.Close(); err != nil {
|
|
return err
|
|
}
|
|
return c.r.Close()
|
|
}
|
|
|
|
func (c *connHalf) SetDeadline(t time.Time) error {
|
|
err1 := c.SetReadDeadline(t)
|
|
err2 := c.SetWriteDeadline(t)
|
|
if err1 != nil {
|
|
return err1
|
|
}
|
|
return err2
|
|
}
|
|
func (c *connHalf) SetReadDeadline(t time.Time) error {
|
|
return c.r.SetReadDeadline(t)
|
|
}
|
|
func (c *connHalf) SetWriteDeadline(t time.Time) error {
|
|
return c.w.SetWriteDeadline(t)
|
|
}
|
|
|
|
func (c *connHalf) SetReadBlock(b bool) error {
|
|
if b {
|
|
return c.r.Block()
|
|
}
|
|
return c.r.Unblock()
|
|
}
|
|
func (c *connHalf) SetWriteBlock(b bool) error {
|
|
if b {
|
|
return c.w.Block()
|
|
}
|
|
return c.w.Unblock()
|
|
}
|