diff --git a/waterutil/tun_ipv4.go b/waterutil/tun_ipv4.go index 6878174..6a6bd36 100644 --- a/waterutil/tun_ipv4.go +++ b/waterutil/tun_ipv4.go @@ -28,11 +28,42 @@ func IPv4Source(packet []byte) net.IP { return net.IPv4(packet[12], packet[13], packet[14], packet[15]) } +func SetIPv4Source(packet []byte, source net.IP) { + copy(packet[12:16], source) +} + func IPv4Destination(packet []byte) net.IP { return net.IPv4(packet[16], packet[17], packet[18], packet[19]) } +func SetIPv4Destination(packet []byte, dest net.IP) { + copy(packet[16:20], dest) +} + func IPv4Payload(packet []byte) []byte { ihl := packet[0] & 0x0F return packet[ihl*4:] } + +// For TCP/UDP +func IPv4SourcePort(packet []byte) uint16 { + payload := IPv4Payload(packet) + return (uint16(payload[0]) << 8) | uint16(payload[1]) +} + +func IPv4DestinationPort(packet []byte) uint16 { + payload := IPv4Payload(packet) + return (uint16(payload[2]) << 8) | uint16(payload[3]) +} + +func SetIPv4SourcePort(packet []byte, port uint16) { + payload := IPv4Payload(packet) + payload[0] = byte(port >> 8) + payload[1] = byte(port & 0xFF) +} + +func SetIPv4DestinationPort(packet []byte, port uint16) { + payload := IPv4Payload(packet) + payload[2] = byte(port >> 8) + payload[3] = byte(port & 0xFF) +}