Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 5 additions & 5 deletions protocol/wireguard/endpoint.go
Original file line number Diff line number Diff line change
Expand Up @@ -92,12 +92,12 @@ func NewEndpoint(ctx context.Context, router adapter.Router, logger log.ContextL
Address: options.Address,
PrivateKey: options.PrivateKey,
ListenPort: options.ListenPort,
ResolvePeer: func(domain string) (netip.Addr, error) {
endpointAddresses, lookupErr := ep.dnsRouter.Lookup(ctx, domain, outboundDialer.(dialer.ResolveDialer).QueryOptions())
if lookupErr != nil {
return netip.Addr{}, lookupErr
ResolvePeer: func(resolveCtx context.Context, domain string, disableCache bool) ([]netip.Addr, error) {
queryOptions := outboundDialer.(dialer.ResolveDialer).QueryOptions()
if disableCache {
queryOptions.DisableCache = true
}
return endpointAddresses[0], nil
return ep.dnsRouter.Lookup(resolveCtx, domain, queryOptions)
},
Peers: common.Map(options.Peers, func(it option.WireGuardPeer) wireguard.PeerOptions {
return wireguard.PeerOptions{
Expand Down
173 changes: 158 additions & 15 deletions transport/wireguard/endpoint.go
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ import (
"os"
"reflect"
"strings"
"sync"
"time"
"unsafe"

Expand All @@ -30,18 +31,30 @@ import (
)

type Endpoint struct {
options EndpointOptions
peers []peerConfig
ipcConf string
allowedAddress []netip.Prefix
tunDevice Device
natDevice NatDevice
device *device.Device
allowedIPs *device.AllowedIPs
pause pause.Manager
pauseCallback *list.Element[pause.Callback]
options EndpointOptions
peers []peerConfig
ipcConf string
allowedAddress []netip.Prefix
tunDevice Device
natDevice NatDevice
device *device.Device
allowedIPs *device.AllowedIPs
pause pause.Manager
pauseCallback *list.Element[pause.Callback]
bind *runtimeBind
ipcSet func(string) error
dnsRefresh chan struct{}
dnsRefreshDone chan struct{}
dnsRefreshCancel context.CancelFunc
}

const (
// wireGuardHandshakeRetryLog is emitted by wireguard-go's expiredRetransmitHandshake.
// wireguard-go does not expose a typed callback for this event.
wireGuardHandshakeRetryLog = "%s - Handshake did not complete after %d seconds, retrying (try %d)"
dnsRefreshInterval = 10 * time.Second
)

func NewEndpoint(options EndpointOptions) (*Endpoint, error) {
if options.PrivateKey == "" {
return nil, E.New("missing private key")
Expand Down Expand Up @@ -145,32 +158,38 @@ func (e *Endpoint) Start(resolve bool) error {
if peer.endpoint.IsValid() || !peer.destination.IsDomain() {
continue
}
destinationAddress, err := e.options.ResolvePeer(peer.destination.Fqdn)
destinationAddresses, err := e.options.ResolvePeer(e.options.Context, peer.destination.Fqdn, false)
if err != nil {
return E.Cause(err, "resolve endpoint domain for peer[", peerIndex, "]: ", peer.destination)
}
destinationAddress, loaded := firstValidAddress(destinationAddresses)
if !loaded {
return E.New("no addresses found for peer[", peerIndex, "]: ", peer.destination)
}
e.peers[peerIndex].endpoint = netip.AddrPortFrom(destinationAddress, peer.destination.Port)
}
} else if resolve {
return nil
}
var bind conn.Bind
var rawBind conn.Bind
wgListener, isWgListener := common.Cast[dialer.WireGuardListener](e.options.Dialer)
if isWgListener {
bind = conn.NewStdNetBind(wgListener.WireGuardControl())
rawBind = conn.NewStdNetBind(wgListener.WireGuardControl())
} else {
var (
isConnect bool
connectAddr netip.AddrPort
reserved [3]uint8
)
if len(e.peers) == 1 && e.peers[0].endpoint.IsValid() {
if len(e.peers) == 1 && e.peers[0].endpoint.IsValid() && !e.peers[0].destination.IsDomain() {
isConnect = true
connectAddr = e.peers[0].endpoint
reserved = e.peers[0].reserved
}
bind = NewClientBind(e.options.Context, e.options.Logger, e.options.Dialer, isConnect, connectAddr, reserved)
rawBind = NewClientBind(e.options.Context, e.options.Logger, e.options.Dialer, isConnect, connectAddr, reserved)
}
bind := &runtimeBind{Bind: rawBind}
e.bind = bind
if isWgListener || len(e.peers) > 1 {
for _, peer := range e.peers {
if peer.reserved != [3]uint8{} {
Expand All @@ -184,6 +203,9 @@ func (e *Endpoint) Start(resolve bool) error {
}
logger := &device.Logger{
Verbosef: func(format string, args ...any) {
if isWireGuardHandshakeRetry(format) {
e.triggerDNSRefresh()
}
e.options.Logger.Debug(fmt.Sprintf(strings.ToLower(format), args...))
},
Errorf: func(format string, args ...any) {
Expand All @@ -209,6 +231,8 @@ func (e *Endpoint) Start(resolve bool) error {
return E.Cause(err, "setup wireguard: \n", ipcConf.String())
}
e.device = wgDevice
e.ipcSet = wgDevice.IpcSet
e.startDNSRefresh()
e.pause = service.FromContext[pause.Manager](e.options.Context)
if e.pause != nil {
e.pauseCallback = e.pause.RegisterCallback(e.onPauseUpdated)
Expand All @@ -232,6 +256,7 @@ func (e *Endpoint) ListenPacket(ctx context.Context, destination M.Socksaddr) (n
}

func (e *Endpoint) Close() error {
e.stopDNSRefresh()
if e.pauseCallback != nil {
e.pause.UnregisterCallback(e.pauseCallback)
e.pauseCallback = nil
Expand All @@ -241,6 +266,7 @@ func (e *Endpoint) Close() error {
e.device.Close()
e.device = nil
}
e.ipcSet = nil
return nil
}

Expand All @@ -267,6 +293,123 @@ func (e *Endpoint) onPauseUpdated(event int) {
}
}

func (e *Endpoint) startDNSRefresh() {
if !common.Any(e.peers, func(peer peerConfig) bool { return peer.destination.IsDomain() }) {
return
}
refreshCtx, cancel := context.WithCancel(e.options.Context)
e.dnsRefresh = make(chan struct{}, 1)
e.dnsRefreshDone = make(chan struct{})
e.dnsRefreshCancel = cancel
go func() {
defer close(e.dnsRefreshDone)
var lastRefresh time.Time
for {
select {
case <-refreshCtx.Done():
return
case <-e.dnsRefresh:
if time.Since(lastRefresh) < dnsRefreshInterval {
continue
}
e.refreshPeerEndpoints(refreshCtx)
lastRefresh = time.Now()
}
}
}()
}

func (e *Endpoint) stopDNSRefresh() {
if e.dnsRefreshCancel == nil {
return
}
e.dnsRefreshCancel()
<-e.dnsRefreshDone
e.dnsRefreshCancel = nil
}

func (e *Endpoint) triggerDNSRefresh() {
if e.dnsRefresh == nil {
return
}
select {
case e.dnsRefresh <- struct{}{}:
default:
}
}

func (e *Endpoint) refreshPeerEndpoints(ctx context.Context) {
if e.ipcSet == nil {
return
}
for peerIndex := range e.peers {
peer := &e.peers[peerIndex]
if !peer.destination.IsDomain() {
continue
}
addresses, err := e.options.ResolvePeer(ctx, peer.destination.Fqdn, true)
if err != nil {
e.options.Logger.Warn(E.Cause(err, "resolve WireGuard peer endpoint: ", peer.destination.Fqdn))
continue
}
newAddress, loaded := firstDifferentAddress(addresses, peer.endpoint.Addr())
if !loaded {
continue
}
oldEndpoint := peer.endpoint
newEndpoint := netip.AddrPortFrom(newAddress, peer.destination.Port)
if peer.reserved != [3]uint8{} && e.bind != nil {
e.bind.SetReservedForEndpoint(newEndpoint, peer.reserved)
}
ipcConf := "public_key=" + peer.publicKeyHex + "\nupdate_only=true\nendpoint=" + newEndpoint.String()
if err = e.ipcSet(ipcConf); err != nil {
e.options.Logger.Warn(E.Cause(err, "update WireGuard peer endpoint: ", peer.destination.Fqdn))
continue
}
peer.endpoint = newEndpoint
e.options.Logger.Info("updated WireGuard peer endpoint for ", peer.destination.Fqdn, ": ", oldEndpoint, " -> ", newEndpoint)
}
}

func firstValidAddress(addresses []netip.Addr) (netip.Addr, bool) {
for _, address := range addresses {
if address.IsValid() {
return address, true
}
}
return netip.Addr{}, false
}

func firstDifferentAddress(addresses []netip.Addr, current netip.Addr) (netip.Addr, bool) {
for _, address := range addresses {
if address.IsValid() && address != current {
return address, true
}
}
return netip.Addr{}, false
}

func isWireGuardHandshakeRetry(format string) bool {
return format == wireGuardHandshakeRetryLog
}

type runtimeBind struct {
conn.Bind
access sync.RWMutex
}

func (b *runtimeBind) Send(buffers [][]byte, endpoint conn.Endpoint, offset int) error {
b.access.RLock()
defer b.access.RUnlock()
return b.Bind.Send(buffers, endpoint, offset)
}

func (b *runtimeBind) SetReservedForEndpoint(destination netip.AddrPort, reserved [3]byte) {
b.access.Lock()
defer b.access.Unlock()
b.Bind.SetReservedForEndpoint(destination, reserved)
}

type peerConfig struct {
destination M.Socksaddr
endpoint netip.AddrPort
Expand Down
2 changes: 1 addition & 1 deletion transport/wireguard/endpoint_options.go
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@ type EndpointOptions struct {
Address []netip.Prefix
PrivateKey string
ListenPort uint16
ResolvePeer func(domain string) (netip.Addr, error)
ResolvePeer func(ctx context.Context, domain string, disableCache bool) ([]netip.Addr, error)
Peers []PeerOptions
Workers int
}
Expand Down
116 changes: 116 additions & 0 deletions transport/wireguard/endpoint_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,116 @@
package wireguard

import (
"context"
"errors"
"net/netip"
"strings"
"sync/atomic"
"testing"
"time"

"github.com/sagernet/sing/common/logger"
M "github.com/sagernet/sing/common/metadata"
"github.com/stretchr/testify/require"
)

func TestWireGuardHandshakeRetryLog(t *testing.T) {
require.True(t, isWireGuardHandshakeRetry(wireGuardHandshakeRetryLog))
require.False(t, isWireGuardHandshakeRetry("%s - Sending handshake initiation"))
require.False(t, isWireGuardHandshakeRetry("%s - Handshake did not complete after %d attempts, giving up"))
}

func TestRefreshPeerEndpoint(t *testing.T) {
oldAddress := netip.MustParseAddr("192.0.2.1")
newAddress := netip.MustParseAddr("192.0.2.2")
var disableCache bool
var ipcConfig string
endpoint := newRefreshTestEndpoint(oldAddress, func(ctx context.Context, domain string, noCache bool) ([]netip.Addr, error) {
require.Equal(t, "peer.example", domain)
disableCache = noCache
return []netip.Addr{oldAddress, newAddress}, nil
})
endpoint.ipcSet = func(config string) error {
ipcConfig = config
return nil
}

endpoint.refreshPeerEndpoints(context.Background())

require.True(t, disableCache)
require.Equal(t, netip.AddrPortFrom(newAddress, 51820), endpoint.peers[0].endpoint)
require.Equal(t, "public_key="+strings.Repeat("ab", 32)+"\nupdate_only=true\nendpoint=192.0.2.2:51820", ipcConfig)
}

func TestRefreshPeerEndpointUnchanged(t *testing.T) {
oldAddress := netip.MustParseAddr("192.0.2.1")
endpoint := newRefreshTestEndpoint(oldAddress, func(ctx context.Context, domain string, noCache bool) ([]netip.Addr, error) {
return []netip.Addr{oldAddress}, nil
})
endpoint.ipcSet = func(config string) error {
t.Fatal("unexpected WireGuard IPC update")
return nil
}

endpoint.refreshPeerEndpoints(context.Background())

require.Equal(t, netip.AddrPortFrom(oldAddress, 51820), endpoint.peers[0].endpoint)
}

func TestRefreshPeerEndpointKeepsOldAddressOnFailure(t *testing.T) {
oldAddress := netip.MustParseAddr("192.0.2.1")
newAddress := netip.MustParseAddr("192.0.2.2")
endpoint := newRefreshTestEndpoint(oldAddress, func(ctx context.Context, domain string, noCache bool) ([]netip.Addr, error) {
return []netip.Addr{newAddress}, nil
})
endpoint.ipcSet = func(config string) error {
return errors.New("update failed")
}

endpoint.refreshPeerEndpoints(context.Background())

require.Equal(t, netip.AddrPortFrom(oldAddress, 51820), endpoint.peers[0].endpoint)
}

func TestDNSRefreshEventsAreCoalescedAndThrottled(t *testing.T) {
oldAddress := netip.MustParseAddr("192.0.2.1")
var lookupCount atomic.Int32
lookupDone := make(chan struct{}, 1)
endpoint := newRefreshTestEndpoint(oldAddress, func(ctx context.Context, domain string, noCache bool) ([]netip.Addr, error) {
lookupCount.Add(1)
lookupDone <- struct{}{}
return []netip.Addr{oldAddress}, nil
})
endpoint.ipcSet = func(config string) error { return nil }
endpoint.startDNSRefresh()
t.Cleanup(endpoint.stopDNSRefresh)

endpoint.triggerDNSRefresh()
endpoint.triggerDNSRefresh()
select {
case <-lookupDone:
case <-time.After(time.Second):
t.Fatal("DNS refresh did not run")
}
endpoint.triggerDNSRefresh()
time.Sleep(50 * time.Millisecond)

require.Equal(t, int32(1), lookupCount.Load())
}

func newRefreshTestEndpoint(oldAddress netip.Addr, resolve func(context.Context, string, bool) ([]netip.Addr, error)) *Endpoint {
return &Endpoint{
options: EndpointOptions{
Context: context.Background(),
Logger: logger.NOP(),
ResolvePeer: resolve,
},
peers: []peerConfig{
{
destination: M.ParseSocksaddrHostPort("peer.example", 51820),
endpoint: netip.AddrPortFrom(oldAddress, 51820),
publicKeyHex: strings.Repeat("ab", 32),
},
},
}
}