diff --git a/examples/main.go b/examples/main.go index 8f17c16..431955e 100644 --- a/examples/main.go +++ b/examples/main.go @@ -5,7 +5,7 @@ import ( "fmt" "log" "time" - + "github.com/Snider/Enchantrix/pkg/crypt" "github.com/Snider/Enchantrix/pkg/crypt/std/chachapoly" "github.com/Snider/Enchantrix/pkg/trix" ) @@ -20,9 +20,9 @@ func main() { // 2. Create a Trix container with the plaintext and attach sigils trixContainer := &trix.Trix{ - Header: map[string]interface{}{}, - Payload: plaintext, - Sigils: []trix.Sigil{&trix.ReverseSigil{}}, + Header: map[string]interface{}{}, + Payload: plaintext, + InSigils: []trix.Sigil{&trix.ReverseSigil{}}, } // 3. Pack the Trix container to apply the sigil transformations @@ -39,7 +39,7 @@ func main() { } trixContainer.Payload = ciphertext // Update the payload with the ciphertext - // 5. Add encryption metadata to the header + // 5. Add encryption metadata and checksum to the header nonce := ciphertext[:24] trixContainer.Header = map[string]interface{}{ "content_type": "application/octet-stream", @@ -47,6 +47,7 @@ func main() { "nonce": base64.StdEncoding.EncodeToString(nonce), "created_at": time.Now().UTC().Format(time.RFC3339), } + trixContainer.ChecksumAlgo = crypt.SHA256 // 6. Encode the .trix container into its binary format @@ -73,7 +74,7 @@ func main() { decodedTrix.Payload = decryptedPayload // 9. Unpack the Trix container to reverse the sigil transformations - decodedTrix.Sigils = trixContainer.Sigils // Re-attach sigils + decodedTrix.InSigils = trixContainer.InSigils // Re-attach sigils if err := decodedTrix.Unpack(); err != nil { log.Fatalf("Failed to unpack trix container: %v", err) } diff --git a/pkg/trix/trix.go b/pkg/trix/trix.go index b2b5dcd..4613a9c 100644 --- a/pkg/trix/trix.go +++ b/pkg/trix/trix.go @@ -7,6 +7,7 @@ import ( "errors" "fmt" "io" + "github.com/Snider/Enchantrix/pkg/crypt" ) const ( @@ -18,6 +19,7 @@ var ( ErrInvalidVersion = errors.New("trix: invalid version") ErrMagicNumberLength = errors.New("trix: magic number must be 4 bytes long") ErrNilSigil = errors.New("trix: sigil cannot be nil") + ErrChecksumMismatch = errors.New("trix: checksum mismatch") ) // Sigil defines the interface for a data transformer. @@ -28,9 +30,11 @@ type Sigil interface { // Trix represents the structure of a .trix file. type Trix struct { - Header map[string]interface{} - Payload []byte - Sigils []Sigil `json:"-"` // Ignore Sigils during JSON marshaling + Header map[string]interface{} + Payload []byte + InSigils []Sigil `json:"-"` // Ignore Sigils during JSON marshaling + OutSigils []Sigil `json:"-"` // Ignore Sigils during JSON marshaling + ChecksumAlgo crypt.HashType `json:"-"` } // Encode serializes a Trix struct into the .trix binary format. @@ -39,6 +43,13 @@ func Encode(trix *Trix, magicNumber string) ([]byte, error) { return nil, ErrMagicNumberLength } + // Calculate and add checksum if an algorithm is specified + if trix.ChecksumAlgo != "" { + checksum := crypt.NewService().Hash(trix.ChecksumAlgo, string(trix.Payload)) + trix.Header["checksum"] = checksum + trix.Header["checksum_algo"] = string(trix.ChecksumAlgo) + } + headerBytes, err := json.Marshal(trix.Header) if err != nil { return nil, err @@ -124,6 +135,18 @@ func Decode(data []byte, magicNumber string) (*Trix, error) { return nil, err } + // Verify checksum if it exists in the header + if checksum, ok := header["checksum"].(string); ok { + algo, ok := header["checksum_algo"].(string) + if !ok { + return nil, errors.New("trix: checksum algorithm not found in header") + } + expectedChecksum := crypt.NewService().Hash(crypt.HashType(algo), string(payload)) + if checksum != expectedChecksum { + return nil, ErrChecksumMismatch + } + } + return &Trix{ Header: header, Payload: payload, @@ -132,7 +155,7 @@ func Decode(data []byte, magicNumber string) (*Trix, error) { // Pack applies the In method of all attached sigils to the payload. func (t *Trix) Pack() error { - for _, sigil := range t.Sigils { + for _, sigil := range t.InSigils { if sigil == nil { return ErrNilSigil } @@ -147,8 +170,12 @@ func (t *Trix) Pack() error { // Unpack applies the Out method of all sigils in reverse order. func (t *Trix) Unpack() error { - for i := len(t.Sigils) - 1; i >= 0; i-- { - sigil := t.Sigils[i] + sigils := t.OutSigils + if len(sigils) == 0 { + sigils = t.InSigils + } + for i := len(sigils) - 1; i >= 0; i-- { + sigil := sigils[i] if sigil == nil { return ErrNilSigil } diff --git a/pkg/trix/trix_test.go b/pkg/trix/trix_test.go index a9acb9f..08578ad 100644 --- a/pkg/trix/trix_test.go +++ b/pkg/trix/trix_test.go @@ -5,7 +5,7 @@ import ( "io" "reflect" "testing" - + "github.com/Snider/Enchantrix/pkg/crypt" "github.com/stretchr/testify/assert" ) @@ -130,9 +130,9 @@ func (s *FailingSigil) Out(data []byte) ([]byte, error) { func TestPackUnpack_Good(t *testing.T) { originalPayload := []byte("hello world") trix := &Trix{ - Header: map[string]interface{}{}, - Payload: originalPayload, - Sigils: []Sigil{&ReverseSigil{}, &ReverseSigil{}}, // Double reverse should be original + Header: map[string]interface{}{}, + Payload: originalPayload, + InSigils: []Sigil{&ReverseSigil{}, &ReverseSigil{}}, // Double reverse should be original } err := trix.Pack() @@ -147,9 +147,9 @@ func TestPackUnpack_Good(t *testing.T) { func TestPackUnpack_Bad(t *testing.T) { expectedErr := errors.New("sigil failed") trix := &Trix{ - Header: map[string]interface{}{}, - Payload: []byte("some data"), - Sigils: []Sigil{&ReverseSigil{}, &FailingSigil{err: expectedErr}}, + Header: map[string]interface{}{}, + Payload: []byte("some data"), + InSigils: []Sigil{&ReverseSigil{}, &FailingSigil{err: expectedErr}}, } err := trix.Pack() @@ -160,9 +160,9 @@ func TestPackUnpack_Bad(t *testing.T) { func TestPackUnpack_Ugly(t *testing.T) { t.Run("NilSigil", func(t *testing.T) { trix := &Trix{ - Header: map[string]interface{}{}, - Payload: []byte("some data"), - Sigils: []Sigil{nil}, + Header: map[string]interface{}{}, + Payload: []byte("some data"), + InSigils: []Sigil{nil}, } err := trix.Pack() @@ -170,3 +170,60 @@ func TestPackUnpack_Ugly(t *testing.T) { assert.Equal(t, ErrNilSigil, err) }) } + +// --- Checksum Tests --- + +func TestChecksum_Good(t *testing.T) { + trix := &Trix{ + Header: map[string]interface{}{}, + Payload: []byte("hello world"), + ChecksumAlgo: crypt.SHA256, + } + encoded, err := Encode(trix, "CHCK") + assert.NoError(t, err) + + decoded, err := Decode(encoded, "CHCK") + assert.NoError(t, err) + assert.Equal(t, trix.Payload, decoded.Payload) +} + +func TestChecksum_Bad(t *testing.T) { + trix := &Trix{ + Header: map[string]interface{}{}, + Payload: []byte("hello world"), + ChecksumAlgo: crypt.SHA256, + } + encoded, err := Encode(trix, "CHCK") + assert.NoError(t, err) + + // Tamper with the payload + encoded[len(encoded)-1] = 0 + + _, err = Decode(encoded, "CHCK") + assert.Error(t, err) + assert.Equal(t, ErrChecksumMismatch, err) +} + +func TestChecksum_Ugly(t *testing.T) { + t.Run("MissingAlgoInHeader", func(t *testing.T) { + trix := &Trix{ + Header: map[string]interface{}{}, + Payload: []byte("hello world"), + ChecksumAlgo: crypt.SHA256, + } + encoded, err := Encode(trix, "UGLY") + assert.NoError(t, err) + + // Manually decode to tamper with the header + decoded, err := Decode(encoded, "UGLY") + assert.NoError(t, err) + delete(decoded.Header, "checksum_algo") + + // Re-encode with the tampered header + tamperedEncoded, err := Encode(decoded, "UGLY") + assert.NoError(t, err) + + _, err = Decode(tamperedEncoded, "UGLY") + assert.Error(t, err) + }) +}