Support AES GCM and snappy compression

5 files changed, 119 insertions, 10 deletions

diff --git a/go.mod b/go.mod
new file mode 100644
index 0000000..6adbd45
--- /dev/null
+++ b/go.mod
@@ -0,0 +1,5 @@
+module code.crute.me/mcrute/go-inform
+
+go 1.18
+
+require github.com/golang/snappy v0.0.4
diff --git a/go.sum b/go.sum
new file mode 100644
index 0000000..74eae48
--- /dev/null
+++ b/go.sum
@@ -0,0 +1,2 @@
+github.com/golang/snappy v0.0.4 h1:yAGX7huGHXlcLOEtBnF4w7FQwA26wojNCwOYAEhLjQM=
+github.com/golang/snappy v0.0.4/go.mod h1:/XxbfmMg8lxefKM7IXC3fBNl/7bRcc72aCRzEWrmP2Q=
diff --git a/inform/codec.go b/inform/codec.go
index e5c9296..8147986 100644
--- a/inform/codec.go
+++ b/inform/codec.go
@@ -5,6 +5,8 @@ import (
         "encoding/binary"
         "errors"
         "io"
+
+        "github.com/golang/snappy"
 )
 
 type Codec struct {
@@ -33,6 +35,7 @@ func (c *Codec) Unmarshal(fp io.Reader) (*InformWrapper, error) {
 
         var dataLen int32
         binary.Read(fp, binary.BigEndian, &dataLen)
+        w.DataLength = dataLen
 
         p := make([]byte, dataLen)
         io.ReadFull(fp, p)
@@ -42,12 +45,21 @@ func (c *Codec) Unmarshal(fp io.Reader) (*InformWrapper, error) {
                 return nil, errors.New("No key found")
         }
 
-        u, err := Decrypt(p, iv, key)
+        u, err := Decrypt(p, iv, key, w)
         if err != nil {
                 return nil, err
         }
 
-        w.Payload = u
+        if w.IsSnappyCompressed() {
+                w.Payload, err = snappy.Decode(nil, u)
+                if err != nil {
+                        return nil, err
+                }
+        } else if w.IsZlibCompressed() {
+                return nil, errors.New("payload is zlib compressed, not supported")
+        } else {
+                w.Payload = u
+        }
 
         return w, nil
 }
diff --git a/inform/crypto.go b/inform/crypto.go
index 90118c1..4e82741 100644
--- a/inform/crypto.go
+++ b/inform/crypto.go
@@ -5,6 +5,7 @@ import (
         "crypto/aes"
         "crypto/cipher"
         "crypto/rand"
+        "encoding/binary"
         "encoding/hex"
         "errors"
 )
@@ -68,7 +69,54 @@ func Encrypt(payload []byte, key string) ([]byte, []byte, error) {
         return ct, iv, nil
 }
 
-func Decrypt(payload, iv []byte, key string) ([]byte, error) {
+func Decrypt(payload, iv []byte, key string, w *InformWrapper) ([]byte, error) {
+        if !w.IsEncrypted() {
+                return nil, errors.New("payload is not encrypted")
+        }
+
+        if w.IsGCMEncrypted() {
+                return decryptGCM(payload, iv, key, w)
+        } else {
+                return decryptCBC(payload, iv, key)
+        }
+
+        return nil, nil
+}
+
+func buildAuthData(w *InformWrapper, iv []byte) []byte {
+        ad := &bytes.Buffer{}
+        binary.Write(ad, binary.BigEndian, int32(PROTOCOL_MAGIC))
+        binary.Write(ad, binary.BigEndian, int32(w.Version))
+        ad.Write(w.MacAddr)
+        binary.Write(ad, binary.BigEndian, int16(w.Flags))
+        ad.Write(iv)
+        binary.Write(ad, binary.BigEndian, int32(w.DataVersion))
+        binary.Write(ad, binary.BigEndian, int32(w.DataLength))
+        return ad.Bytes()
+}
+
+func decryptGCM(payload, iv []byte, key string, w *InformWrapper) ([]byte, error) {
+        block, err := decodeHexKey(key)
+        if err != nil {
+                return nil, err
+        }
+
+        mode, err := cipher.NewGCMWithNonceSize(block, 16)
+        if err != nil {
+                return nil, err
+        }
+
+        _, err = mode.Open(payload[:0], iv, payload, buildAuthData(w, iv))
+        if err != nil {
+                return nil, err
+        }
+
+        // The last block always seems to be garbage, maybe it's padding or
+        // something else. I have not looked carefully at it.
+        return payload[:len(payload)-aes.BlockSize], nil
+}
+
+func decryptCBC(payload, iv []byte, key string) ([]byte, error) {
         b := make([]byte, len(payload))
 
         block, err := decodeHexKey(key)
diff --git a/inform/inform.go b/inform/inform.go
index ac3b57d..59f969c 100644
--- a/inform/inform.go
+++ b/inform/inform.go
@@ -11,8 +11,10 @@ const (
         INFORM_VERSION int32 = 0
         DATA_VERSION   int32 = 1
 
-        ENCRYPTED_FLAG  = 1
-        COMPRESSED_FLAG = 2
+        ENCRYPTED_FLAG         = 1
+        ZLIB_COMPRESSED_FLAG   = 2
+        SNAPPY_COMPRESSED_FLAG = 4
+        AES_GCM_FLAG           = 8
 )
 
 // Wrapper around an inform message, serializes directly into the wire
@@ -22,6 +24,7 @@ type InformWrapper struct {
         MacAddr     []byte
         Flags       int16
         DataVersion int32
+        DataLength  int32
         Payload     []byte
 }
 
@@ -86,7 +89,18 @@ func (i *InformWrapper) String() string {
         fmt.Fprintf(b, "Mac Addr:     \t%s\n", i.FormattedMac())
         fmt.Fprintf(b, "Flags:        \t%d\n", i.Flags)
         fmt.Fprintf(b, " Encrypted:   \t%t\n", i.IsEncrypted())
-        fmt.Fprintf(b, " Compressed:  \t%t\n", i.IsCompressed())
+        fmt.Fprintf(b, " Compressed:  \t%t", i.IsCompressed())
+        if i.IsCompressed() {
+                if i.IsZlibCompressed() {
+                        fmt.Fprintf(b, " (zlib)\n")
+                } else if i.IsSnappyCompressed() {
+                        fmt.Fprintf(b, " (snappy)\n")
+                } else {
+                        fmt.Fprintf(b, " (unknown)\n")
+                }
+        } else {
+                fmt.Fprintf(b, "\n")
+        }
         fmt.Fprintf(b, "Data Version: \t%d\n", i.DataVersion)
         fmt.Fprintf(b, "Payload:      \t%q\n", i.Payload)
 
@@ -105,14 +119,42 @@ func (i *InformWrapper) SetEncrypted(e bool) {
         }
 }
 
+func (i *InformWrapper) IsGCMEncrypted() bool {
+        return i.Flags&AES_GCM_FLAG != 0
+}
+
+func (i *InformWrapper) SetGCMEncrypted(e bool) {
+        if e {
+                i.Flags |= AES_GCM_FLAG
+        } else {
+                i.Flags &= AES_GCM_FLAG
+        }
+}
+
 func (i *InformWrapper) IsCompressed() bool {
-        return i.Flags&COMPRESSED_FLAG != 0
+        return i.IsZlibCompressed() || i.IsSnappyCompressed()
+}
+
+func (i *InformWrapper) IsZlibCompressed() bool {
+        return i.Flags&ZLIB_COMPRESSED_FLAG != 0
+}
+
+func (i *InformWrapper) IsSnappyCompressed() bool {
+        return i.Flags&SNAPPY_COMPRESSED_FLAG != 0
+}
+
+func (i *InformWrapper) SetSnappyCompressed(c bool) {
+        if c {
+                i.Flags |= SNAPPY_COMPRESSED_FLAG
+        } else {
+                i.Flags &= SNAPPY_COMPRESSED_FLAG
+        }
 }
 
-func (i *InformWrapper) SetCompressed(c bool) {
+func (i *InformWrapper) SetZlibCompressed(c bool) {
         if c {
-                i.Flags |= COMPRESSED_FLAG
+                i.Flags |= ZLIB_COMPRESSED_FLAG
         } else {
-                i.Flags &= COMPRESSED_FLAG
+                i.Flags &= ZLIB_COMPRESSED_FLAG
         }
 }