From c074f56200edc096891551f28b4cf46f76073d1b Mon Sep 17 00:00:00 2001
From: Vishvananda Ishaya <vishvananda@gmail.com>
Date: Mon, 15 Sep 2014 17:04:48 -0700
Subject: [PATCH] Add support for XfrmState Encapsulation

---
 xfrm_state.go            | 44 ++++++++++++++++++++++++++++++++--------
 xfrm_state_linux.go      | 38 ++++++++++++++++++++++++++++++++++
 xfrm_state_linux_test.go | 29 ++++++++++++++++++++++++++
 3 files changed, 103 insertions(+), 8 deletions(-)

diff --git a/xfrm_state.go b/xfrm_state.go
index 550506d..9eaad09 100644
--- a/xfrm_state.go
+++ b/xfrm_state.go
@@ -11,16 +11,44 @@ type XfrmStateAlgo struct {
 	TruncateLen int // Auth only
 }
 
+// EncapType is an enum representing an ipsec template direction.
+type EncapType uint8
+
+const (
+	_                                    = iota
+	XFRM_ENCAP_ESPINUDP_NONIKE EncapType = iota
+	XFRM_ENCAP_ESPINUDP        EncapType = iota
+)
+
+func (e EncapType) String() string {
+	switch e {
+	case XFRM_ENCAP_ESPINUDP_NONIKE:
+		return "espinudp-nonike"
+	case XFRM_ENCAP_ESPINUDP:
+		return "espinudp"
+	}
+	return "unknown"
+}
+
+// XfrmEncap represents the encapsulation to use for the ipsec encryption.
+type XfrmStateEncap struct {
+	Type            EncapType
+	SrcPort         int
+	DstPort         int
+	OriginalAddress net.IP
+}
+
 // XfrmState represents the state of an ipsec policy. It optionally
 // contains an XfrmStateAlgo for encryption and one for authentication.
 type XfrmState struct {
-	Dst   net.IP
-	Src   net.IP
-	Proto Proto
-	Mode  Mode
-	Spi   int
-	Reqid int
+	Dst          net.IP
+	Src          net.IP
+	Proto        Proto
+	Mode         Mode
+	Spi          int
+	Reqid        int
 	ReplayWindow int
-	Auth  *XfrmStateAlgo
-	Crypt *XfrmStateAlgo
+	Auth         *XfrmStateAlgo
+	Crypt        *XfrmStateAlgo
+	Encap        *XfrmStateEncap
 }
diff --git a/xfrm_state_linux.go b/xfrm_state_linux.go
index a56e167..c93d690 100644
--- a/xfrm_state_linux.go
+++ b/xfrm_state_linux.go
@@ -12,6 +12,7 @@ const (
 	SizeofXfrmUsersaInfo = 0xe0
 	SizeofXfrmAlgo       = 0x44
 	SizeofXfrmAlgoAuth   = 0x48
+	SizeofXfrmEncapTmpl  = 0x18
 )
 
 // struct xfrm_usersa_id {
@@ -201,6 +202,26 @@ func (msg *XfrmAlgoAuth) Serialize() []byte {
 //   xfrm_address_t  encap_oa;
 // };
 
+type XfrmEncapTmpl struct {
+	EncapType  uint16
+	EncapSport uint16 // big endian
+	EncapDport uint16 // big endian
+	Pad        [2]byte
+	EncapOa    XfrmAddress
+}
+
+func (msg *XfrmEncapTmpl) Len() int {
+	return SizeofXfrmEncapTmpl
+}
+
+func DeserializeXfrmEncapTmpl(b []byte) *XfrmEncapTmpl {
+	return (*XfrmEncapTmpl)(unsafe.Pointer(&b[0:SizeofXfrmEncapTmpl][0]))
+}
+
+func (msg *XfrmEncapTmpl) Serialize() []byte {
+	return (*(*[SizeofXfrmEncapTmpl]byte)(unsafe.Pointer(msg)))[:]
+}
+
 func writeStateAlgo(a *XfrmStateAlgo) []byte {
 	algo := XfrmAlgo{
 		AlgKeyLen: uint32(len(a.Key) * 8),
@@ -260,6 +281,16 @@ func XfrmStateAdd(state *XfrmState) error {
 		out := newRtAttr(XFRMA_ALG_CRYPT, writeStateAlgo(state.Crypt))
 		req.AddData(out)
 	}
+	if state.Encap != nil {
+		encapData := make([]byte, SizeofXfrmEncapTmpl)
+		encap := DeserializeXfrmEncapTmpl(encapData)
+		encap.EncapType = uint16(state.Encap.Type)
+		encap.EncapSport = swap16(uint16(state.Encap.SrcPort))
+		encap.EncapDport = swap16(uint16(state.Encap.DstPort))
+		encap.EncapOa.FromIP(state.Encap.OriginalAddress)
+		out := newRtAttr(XFRMA_ENCAP, encapData)
+		req.AddData(out)
+	}
 
 	_, err := req.Execute(syscall.NETLINK_XFRM, 0)
 	return err
@@ -349,6 +380,13 @@ func XfrmStateList(family int) ([]XfrmState, error) {
 				state.Auth.Name = bytesToString(algo.AlgName[:])
 				state.Auth.Key = algo.AlgKey
 				state.Auth.TruncateLen = int(algo.AlgTruncLen)
+			case XFRMA_ENCAP:
+				encap := DeserializeXfrmEncapTmpl(attr.Value[:])
+				state.Encap = new(XfrmStateEncap)
+				state.Encap.Type = EncapType(encap.EncapType)
+				state.Encap.SrcPort = int(swap16(encap.EncapSport))
+				state.Encap.DstPort = int(swap16(encap.EncapDport))
+				state.Encap.OriginalAddress = encap.EncapOa.ToIP()
 			}
 
 		}
diff --git a/xfrm_state_linux_test.go b/xfrm_state_linux_test.go
index dbe4982..b8e8fa8 100644
--- a/xfrm_state_linux_test.go
+++ b/xfrm_state_linux_test.go
@@ -176,3 +176,32 @@ func TestXfrmAlgoAuthDeserializeSerialize(t *testing.T) {
 	msg := DeserializeXfrmAlgoAuth(orig)
 	testDeserializeSerialize(t, orig, safemsg, msg)
 }
+
+func (msg *XfrmEncapTmpl) write(b []byte) {
+	native := nativeEndian()
+	native.PutUint16(b[0:2], msg.EncapType)
+	native.PutUint16(b[2:4], msg.EncapSport)
+	native.PutUint16(b[4:6], msg.EncapDport)
+	copy(b[6:8], msg.Pad[:])
+	msg.EncapOa.write(b[8:SizeofXfrmAddress])
+}
+
+func (msg *XfrmEncapTmpl) serializeSafe() []byte {
+	b := make([]byte, SizeofXfrmEncapTmpl)
+	msg.write(b)
+	return b
+}
+
+func deserializeXfrmEncapTmplSafe(b []byte) *XfrmEncapTmpl {
+	var msg = XfrmEncapTmpl{}
+	binary.Read(bytes.NewReader(b[0:SizeofXfrmEncapTmpl]), nativeEndian(), &msg)
+	return &msg
+}
+
+func TestXfrmEncapTmplDeserializeSerialize(t *testing.T) {
+	var orig = make([]byte, SizeofXfrmEncapTmpl)
+	rand.Read(orig)
+	safemsg := deserializeXfrmEncapTmplSafe(orig)
+	msg := DeserializeXfrmEncapTmpl(orig)
+	testDeserializeSerialize(t, orig, safemsg, msg)
+}