add port accessors and modifiers for TCP/UDP

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)
+}