Skip to content

Commit

Permalink
Add winiphlpapi
Browse files Browse the repository at this point in the history
  • Loading branch information
nekohasekai committed Jan 30, 2025
1 parent d9f6eb1 commit 92f3c4c
Show file tree
Hide file tree
Showing 6 changed files with 786 additions and 4 deletions.
6 changes: 2 additions & 4 deletions common/windnsapi/dnsapi_test.go
Original file line number Diff line number Diff line change
@@ -1,16 +1,14 @@
//go:build windows

package windnsapi

import (
"runtime"
"testing"

"github.com/stretchr/testify/require"
)

func TestDNSAPI(t *testing.T) {
if runtime.GOOS != "windows" {
t.SkipNow()
}
t.Parallel()
require.NoError(t, FlushResolverCache())
}
218 changes: 218 additions & 0 deletions common/winiphlpapi/helper.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,218 @@
//go:build windows

package winiphlpapi

import (
"context"
"encoding/binary"
"net"
"net/netip"
"os"
"time"
"unsafe"

E "github.com/sagernet/sing/common/exceptions"
M "github.com/sagernet/sing/common/metadata"
)

func LoadEStats() error {
err := modiphlpapi.Load()
if err != nil {
return err
}
err = procGetTcpTable.Find()
if err != nil {
return err
}
err = procGetTcp6Table.Find()
if err != nil {
return err
}
err = procGetPerTcp6ConnectionEStats.Find()
if err != nil {
return err
}
err = procGetPerTcp6ConnectionEStats.Find()
if err != nil {
return err
}
err = procSetPerTcpConnectionEStats.Find()
if err != nil {
return err
}
err = procSetPerTcp6ConnectionEStats.Find()
if err != nil {
return err
}
return nil
}

func LoadExtendedTable() error {
err := modiphlpapi.Load()
if err != nil {
return err
}
err = procGetExtendedTcpTable.Find()
if err != nil {
return err
}
err = procGetExtendedUdpTable.Find()
if err != nil {
return err
}
return nil
}

func FindTCPPid(source netip.AddrPort, destination netip.AddrPort) (uint32, error) {
if source.Addr().Is4() {
tcpTable, err := GetExtendedTcpTable()
if err != nil {
return 0, err
}
for _, row := range tcpTable {
if source == netip.AddrPortFrom(DwordToAddr(row.DwLocalAddr), DwordToPort(row.DwLocalPort)) ||
destination == netip.AddrPortFrom(DwordToAddr(row.DwRemoteAddr), DwordToPort(row.DwRemotePort)) {
return row.DwOwningPid, nil
}
}
} else {
tcpTable, err := GetExtendedTcp6Table()
if err != nil {
return 0, err
}
for _, row := range tcpTable {
if source == netip.AddrPortFrom(netip.AddrFrom16(row.UcLocalAddr), DwordToPort(row.DwLocalPort)) ||
destination == netip.AddrPortFrom(netip.AddrFrom16(row.UcRemoteAddr), DwordToPort(row.DwRemotePort)) {
return row.DwOwningPid, nil
}
}
}
return 0, E.New("process not found for ", source, "-", destination)
}

func FindUDPPid(source netip.AddrPort) (uint32, error) {
if source.Addr().Is4() {
udpTable, err := GetExtendedUdpTable()
if err != nil {
return 0, err
}
for _, row := range udpTable {
if source == netip.AddrPortFrom(DwordToAddr(row.DwLocalAddr), DwordToPort(row.DwLocalPort)) {
return row.DwOwningPid, nil
}
}
} else {
udpTable, err := GetExtendedUdp6Table()
if err != nil {
return 0, err
}
for _, row := range udpTable {
if source == netip.AddrPortFrom(netip.AddrFrom16(row.UcLocalAddr), DwordToPort(row.DwLocalPort)) {
return row.DwOwningPid, nil
}
}
}
return 0, E.New("process not found for ", source)
}

func WriteAndWaitAck(ctx context.Context, conn net.Conn, payload []byte) error {
source := M.AddrPortFromNet(conn.LocalAddr())
destination := M.AddrPortFromNet(conn.RemoteAddr())
if source.Addr().Is4() {
tcpTable, err := GetTcpTable()
if err != nil {
return err
}
var tcpRow *MibTcpRow
for _, row := range tcpTable {
if source == netip.AddrPortFrom(DwordToAddr(row.DwLocalAddr), DwordToPort(row.DwLocalPort)) ||
destination == netip.AddrPortFrom(DwordToAddr(row.DwRemoteAddr), DwordToPort(row.DwRemotePort)) {
tcpRow = &row
break
}
}
if tcpRow == nil {
return E.New("row not found for: ", source)
}
err = SetPerTcpConnectionEStatsSendBuffer(tcpRow, &TcpEstatsSendBuffRwV0{
EnableCollection: true,
})
if err != nil {
return os.NewSyscallError("SetPerTcpConnectionEStatsSendBufferV0", err)
}
defer SetPerTcpConnectionEStatsSendBuffer(tcpRow, &TcpEstatsSendBuffRwV0{
EnableCollection: false,
})
_, err = conn.Write(payload)
if err != nil {
return err
}
for {
select {
case <-ctx.Done():
return ctx.Err()
default:
}
eStstsSendBuffer, err := GetPerTcpConnectionEStatsSendBuffer(tcpRow)
if err != nil {
return err
}
if eStstsSendBuffer.CurRetxQueue == 0 {
return nil
}
time.Sleep(10 * time.Millisecond)
}
} else {
tcpTable, err := GetTcp6Table()
if err != nil {
return err
}
var tcpRow *MibTcp6Row
for _, row := range tcpTable {
if source == netip.AddrPortFrom(netip.AddrFrom16(row.LocalAddr), DwordToPort(row.LocalPort)) ||
destination == netip.AddrPortFrom(netip.AddrFrom16(row.RemoteAddr), DwordToPort(row.RemotePort)) {
tcpRow = &row
break
}
}
if tcpRow == nil {
return E.New("row not found for: ", source)
}
err = SetPerTcp6ConnectionEStatsSendBuffer(tcpRow, &TcpEstatsSendBuffRwV0{
EnableCollection: true,
})
if err != nil {
return os.NewSyscallError("SetPerTcpConnectionEStatsSendBufferV0", err)
}
defer SetPerTcp6ConnectionEStatsSendBuffer(tcpRow, &TcpEstatsSendBuffRwV0{
EnableCollection: false,
})
_, err = conn.Write(payload)
if err != nil {
return err
}
for {
select {
case <-ctx.Done():
return ctx.Err()
default:
}
eStstsSendBuffer, err := GetPerTcp6ConnectionEStatsSendBuffer(tcpRow)
if err != nil {
return err
}
if eStstsSendBuffer.CurRetxQueue == 0 {
return nil
}
time.Sleep(10 * time.Millisecond)
}
}
}

func DwordToAddr(addr uint32) netip.Addr {
return netip.AddrFrom4(*(*[4]byte)(unsafe.Pointer(&addr)))
}

func DwordToPort(dword uint32) uint16 {
return binary.BigEndian.Uint16((*[4]byte)(unsafe.Pointer(&dword))[:])
}
Loading

0 comments on commit 92f3c4c

Please sign in to comment.