Skip to content

Instantly share code, notes, and snippets.

@Zenithar

Zenithar/idmask.go

Last active October 11, 2022 12:57
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save Zenithar/c1ee4db2cea88c60a3f289bb882d9818 to your computer and use it in GitHub Desktop.
Save Zenithar/c1ee4db2cea88c60a3f289bb882d9818 to your computer and use it in GitHub Desktop.
Download ZIP
ID Mask - Inspired from https://github.com/patrickfav/id-mask
Raw
idmask.go
// Copyright (C) 2020-2023 by Thibault NORMAND <me+oss@zenithar.org>
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
//
// This Source Code Form is "Incompatible With Secondary Licenses", as
// defined by the Mozilla Public License, version 2.0.
package idmask
import (
"crypto/aes"
"crypto/cipher"
"crypto/hmac"
"crypto/sha256"
"crypto/subtle"
"errors"
"fmt"
"io"
"golang.org/x/crypto/blake2b"
"golang.org/x/crypto/chacha20"
"golang.org/x/crypto/hkdf"
)
type IdentifierMask interface {
Mask(plainId []byte) ([]byte, error)
UnMask(encryptedId []byte) ([]byte, error)
}
// -----------------------------------------------------------------------------
func Random8(r io.Reader, key []byte) IdentifierMask {
// Compute the key
hm := hmac.New(sha256.New, key)
hm.Write([]byte{0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01})
dk := hm.Sum(nil)
var key16 [16]byte
copy(key16[:], dk[:16])
return &idMask8{
key: key16,
randReader: r,
}
}
func Determistic16(key []byte) IdentifierMask {
return &idMask16{
key: key,
randomize: false,
}
}
func Random16(r io.Reader, key []byte) IdentifierMask {
return &idMask16{
key: key,
randomize: true,
randReader: r,
}
}
func Determistic16Modern(key []byte) IdentifierMask {
return &idMask16Modern{
key: key,
randomize: false,
}
}
func Random16Modern(r io.Reader, key []byte) IdentifierMask {
return &idMask16Modern{
key: key,
randomize: true,
randReader: r,
}
}
// -----------------------------------------------------------------------------
const (
entropyLength = 16
macLength = 16
encryptionKeyLength = 16
ivLength = 16
authenticationKeyLength = 32
)
type idMask16 struct {
key []byte
randomize bool
randReader io.Reader
}
func (im *idMask16) Mask(identifier []byte) ([]byte, error) {
// Check arguments
if len(identifier) != 16 {
return nil, errors.New("plain identifier must be 16 bytes long for this masking strategy")
}
// Generate a random entropy
var entropy [entropyLength]byte
if im.randomize {
if _, err := io.ReadFull(im.randReader, entropy[:]); err != nil {
return nil, fmt.Errorf("unable to generate initial entropy: %w", err)
}
}
// Key derivation form secret key and random entropy
var dk [encryptionKeyLength + ivLength + authenticationKeyLength]byte
hr := hkdf.New(sha256.New, im.key[:], entropy[:], []byte("idmask:key:generation:v1"))
if _, err := io.ReadFull(hr, dk[:]); err != nil {
return nil, fmt.Errorf("unable to derive keys: %w", err)
}
// Extract keys
var (
ek = dk[:encryptionKeyLength]
iv = dk[encryptionKeyLength : encryptionKeyLength+ivLength]
ak = dk[encryptionKeyLength+ivLength : encryptionKeyLength+ivLength+authenticationKeyLength]
)
// Initialize block cipher
block, err := aes.NewCipher(ek)
if err != nil {
return nil, fmt.Errorf("unable to initialize AES-128 block cipher: %w", err)
}
// Initialize AES-CBC mode
encryptor := cipher.NewCBCEncrypter(block, iv)
// Apply XOR to distribute entropy to plainId
var plainId [16]byte
for i, b := range identifier {
plainId[i] = b ^ entropy[i]
}
// Encrypt the identifier
var encryptedId [16]byte
encryptor.CryptBlocks(encryptedId[:], plainId[:])
// Compute MAC
hm := hmac.New(sha256.New, ak)
hm.Write([]byte("idmask:mac"))
hm.Write(iv)
if im.randomize {
hm.Write(entropy[:])
}
hm.Write(encryptedId[:])
mac := hm.Sum(nil)
// Prepare output
var final []byte
if im.randomize {
final = make([]byte, 0, entropyLength+16+macLength)
final = append(final[:], entropy[:]...)
} else {
final = make([]byte, 0, 16+macLength)
}
final = append(final[:], encryptedId[:]...)
final = append(final[:], mac[:macLength]...)
// Return final masked identifier
return final, nil
}
func (im *idMask16) UnMask(encodedId []byte) ([]byte, error) {
// Key derivation form secret key and random entropy
var (
dk [encryptionKeyLength + ivLength + authenticationKeyLength]byte
zero [entropyLength]byte
hr io.Reader
)
if im.randomize {
hr = hkdf.New(sha256.New, im.key[:], encodedId[:entropyLength], []byte("idmask:key:generation:v1"))
} else {
hr = hkdf.New(sha256.New, im.key[:], zero[:], []byte("idmask:key:generation:v1"))
}
if _, err := io.ReadFull(hr, dk[:]); err != nil {
return nil, fmt.Errorf("unable to derive keys: %w", err)
}
// Extract keys
var (
ek = dk[:encryptionKeyLength]
iv = dk[encryptionKeyLength : encryptionKeyLength+ivLength]
ak = dk[encryptionKeyLength+ivLength : encryptionKeyLength+ivLength+authenticationKeyLength]
)
offset := 0
if im.randomize {
offset = entropyLength
}
// Compute MAC
hm := hmac.New(sha256.New, ak)
hm.Write([]byte("idmask:mac"))
hm.Write(iv)
if im.randomize {
hm.Write(encodedId[:entropyLength])
}
hm.Write(encodedId[offset : 16+offset])
mac := hm.Sum(nil)
// Validate MAC
if subtle.ConstantTimeCompare(mac[:macLength], encodedId[16+offset:]) != 1 {
return nil, errors.New("identifier integrity check failed")
}
// Initialize block cipher
block, err := aes.NewCipher(ek)
if err != nil {
return nil, fmt.Errorf("unable to initialize AES-128 block cipher: %w", err)
}
// Initialize AES-CBC mode
decrypter := cipher.NewCBCDecrypter(block, iv)
// Decrypt the identifier
var plainId [16]byte
decrypter.CryptBlocks(plainId[:], encodedId[offset:16+offset])
// Apply XOR to reverse entropy distribution
if im.randomize {
for i, b := range plainId {
plainId[i] = b ^ encodedId[i]
}
}
return plainId[:], nil
}
type idMask16Modern struct {
key []byte
randomize bool
randReader io.Reader
}
func (im *idMask16Modern) Mask(identifier []byte) ([]byte, error) {
// Check arguments
if len(identifier) != 16 {
return nil, errors.New("plain identifier must be 16 bytes long for this masking strategy")
}
// Generate a random entropy
var entropy [entropyLength]byte
if im.randomize {
if _, err := io.ReadFull(im.randReader, entropy[:]); err != nil {
return nil, fmt.Errorf("unable to generate initial entropy: %w", err)
}
}
// Key derivation form secret key and random entropy
var dk []byte
h, err := blake2b.New(encryptionKeyLength+ivLength+authenticationKeyLength, im.key)
if err != nil {
return nil, fmt.Errorf("unable to derive keys: %w", err)
}
h.Write([]byte("idmask:key:generation:v1"))
h.Write(entropy[:])
dk = h.Sum(nil)
// Extract keys
var (
ek = dk[:32]
iv = dk[32:44]
ak = dk[44:]
)
// Initialize block cipher
cipher, err := chacha20.NewUnauthenticatedCipher(ek, iv)
if err != nil {
return nil, fmt.Errorf("unable to initialize ChaCha20 block cipher: %w", err)
}
// Apply XOR to distribute entropy to plainId
var plainId [16]byte
for i, b := range identifier {
plainId[i] = b ^ entropy[i]
}
// Encrypt the identifier
var encryptedId [16]byte
cipher.XORKeyStream(encryptedId[:], plainId[:])
// Compute MAC
hm, err := blake2b.New(macLength, ak)
if err != nil {
return nil, fmt.Errorf("unable to initialize blake2b MAC: %w", err)
}
hm.Write([]byte("idmask:mac"))
hm.Write(iv)
if im.randomize {
hm.Write(entropy[:])
}
hm.Write(encryptedId[:])
mac := hm.Sum(nil)
// Prepare output
var final []byte
if im.randomize {
final = make([]byte, 0, entropyLength+16+macLength)
final = append(final[:], entropy[:]...)
} else {
final = make([]byte, 0, 16+macLength)
}
final = append(final[:], encryptedId[:]...)
final = append(final[:], mac[:macLength]...)
// Return final masked identifier
return final, nil
}
func (im *idMask16Modern) UnMask(encodedId []byte) ([]byte, error) {
// Key derivation form secret key and random entropy
var (
zero [entropyLength]byte
)
// Key derivation form secret key and random entropy
var dk []byte
h, err := blake2b.New(encryptionKeyLength+ivLength+authenticationKeyLength, im.key)
if err != nil {
return nil, fmt.Errorf("unable to derive keys: %w", err)
}
h.Write([]byte("idmask:key:generation:v1"))
if im.randomize {
h.Write(encodedId[:entropyLength])
} else {
h.Write(zero[:])
}
dk = h.Sum(nil)
// Extract keys
var (
ek = dk[:32]
iv = dk[32:44]
ak = dk[44:]
)
offset := 0
if im.randomize {
offset = entropyLength
}
// Compute MAC
hm, err := blake2b.New(macLength, ak)
if err != nil {
return nil, fmt.Errorf("unable to initialize blake2b MAC: %w", err)
}
hm.Write([]byte("idmask:mac"))
hm.Write(iv)
if im.randomize {
hm.Write(encodedId[:entropyLength])
}
hm.Write(encodedId[offset : 16+offset])
mac := hm.Sum(nil)
// Validate MAC
if subtle.ConstantTimeCompare(mac[:macLength], encodedId[16+offset:]) != 1 {
return nil, errors.New("identifier integrity check failed")
}
// Initialize block cipher
cipher, err := chacha20.NewUnauthenticatedCipher(ek, iv)
if err != nil {
return nil, fmt.Errorf("unable to initialize ChaCha20 block cipher: %w", err)
}
// Decrypt the identifier
var plainId [16]byte
cipher.XORKeyStream(plainId[:], encodedId[offset:16+offset])
// Apply XOR to reverse entropy distribution
if im.randomize {
for i, b := range plainId {
plainId[i] = b ^ encodedId[i]
}
}
return plainId[:], nil
}
// -----------------------------------------------------------------------------
type idMask8 struct {
key [16]byte
randReader io.Reader
}
func (im *idMask8) Mask(plainId []byte) ([]byte, error) {
// Check arguments
if len(plainId) != 8 {
return nil, errors.New("plain identifier must be 8 bytes long for this masking strategy")
}
// Generate a random entropy
var entropy [8]byte
if _, err := io.ReadFull(im.randReader, entropy[:]); err != nil {
return nil, fmt.Errorf("unable to generate initial entropy: %w", err)
}
// Initialize block cipher
block, err := aes.NewCipher(im.key[:])
if err != nil {
return nil, fmt.Errorf("unable to initialize AES-128 block cipher: %w", err)
}
var encryptedId [16]byte
block.Encrypt(encryptedId[:], append(entropy[:], plainId...))
// Return final masked identifier
return append(entropy[:], encryptedId[:]...), nil
}
func (im *idMask8) UnMask(encodedId []byte) ([]byte, error) {
// Initialize block cipher
block, err := aes.NewCipher(im.key[:])
if err != nil {
return nil, fmt.Errorf("unable to initialize AES-128 block cipher: %w", err)
}
var plainId [16]byte
block.Decrypt(plainId[:], encodedId[8:])
// Ensure not tampered content
if subtle.ConstantTimeCompare(encodedId[:8], plainId[:8]) != 1 {
return nil, errors.New("invalid or tampered identifier")
}
// Return final masked identifier
return plainId[8:], nil
}
Raw
idmask_test.go
// Copyright (C) 2020-2023 by Thibault NORMAND <me+oss@zenithar.org>
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
//
// This Source Code Form is "Incompatible With Secondary Licenses", as
// defined by the Mozilla Public License, version 2.0.
package idmask
import (
"bytes"
crand "crypto/rand"
"encoding/base64"
"encoding/binary"
"fmt"
"testing"
"github.com/google/uuid"
"github.com/sony/sonyflake"
)
func Test16(t *testing.T) {
t.Run("ClassicDeterministic", func(t *testing.T) {
idm := Determistic16([]byte("829bph7qn7KXc27NgQN1tAWuinYo0C2Q97bVk9sDN5SilqxHTgcBTMxQraMdbTD"))
testUUIDIdentifier(t, idm)
})
t.Run("ClassicRandomized", func(t *testing.T) {
idm := Random16(crand.Reader, []byte("829bph7qn7KXc27NgQN1tAWuinYo0C2Q97bVk9sDN5SilqxHTgcBTMxQraMdbTD"))
testUUIDIdentifier(t, idm)
})
t.Run("ModernDetermistic", func(t *testing.T) {
idm := Determistic16Modern([]byte("829bph7qn7KXc27NgQN1tAWuinYo0C2Q97bVk9sDN5SilqxHTgcBTMxQraMdbTD"))
testUUIDIdentifier(t, idm)
})
t.Run("ModernRandomized", func(t *testing.T) {
idm := Random16Modern(crand.Reader, []byte("829bph7qn7KXc27NgQN1tAWuinYo0C2Q97bVk9sDN5SilqxHTgcBTMxQraMdbTD"))
testUUIDIdentifier(t, idm)
})
}
func Benchmark16(b *testing.B) {
b.ReportAllocs()
b.Run("ClassicDeterministic", func(b *testing.B) {
idm := Determistic16([]byte("829bph7qn7KXc27NgQN1tAWuinYo0C2Q97bVk9sDN5SilqxHTgcBTMxQraMdbTD"))
benchmark16(b, idm)
})
b.Run("ClassicRandomized", func(b *testing.B) {
idm := Random16(crand.Reader, []byte("829bph7qn7KXc27NgQN1tAWuinYo0C2Q97bVk9sDN5SilqxHTgcBTMxQraMdbTD"))
benchmark16(b, idm)
})
b.Run("ModernDetermistic", func(b *testing.B) {
idm := Determistic16Modern([]byte("829bph7qn7KXc27NgQN1tAWuinYo0C2Q97bVk9sDN5SilqxHTgcBTMxQraMdbTD"))
benchmark16(b, idm)
})
b.Run("ModernRandomized", func(b *testing.B) {
idm := Random16Modern(crand.Reader, []byte("829bph7qn7KXc27NgQN1tAWuinYo0C2Q97bVk9sDN5SilqxHTgcBTMxQraMdbTD"))
benchmark16(b, idm)
})
}
func benchmark16(b *testing.B, idm IdentifierMask) {
b.Run("Mask", func(b *testing.B) {
b.ReportAllocs()
uid := uuid.New()
raw, _ := uid.MarshalBinary()
for n := 0; n < b.N; n++ {
b.SetBytes(int64(len(raw)))
idm.Mask(raw)
}
})
b.Run("UnMask", func(b *testing.B) {
b.ReportAllocs()
uid := uuid.New()
raw, _ := uid.MarshalBinary()
out, _ := idm.Mask(raw)
for n := 0; n < b.N; n++ {
b.SetBytes(int64(len(out)))
idm.UnMask(out)
}
})
}
func testUUIDIdentifier(t *testing.T, idm IdentifierMask) {
uid := uuid.New()
raw, _ := uid.MarshalBinary()
out, err := idm.Mask(raw)
if err != nil {
t.Fatal(err)
}
fmt.Printf("%s => %s\n", uid, base64.RawURLEncoding.EncodeToString(out))
id, err := idm.UnMask(out)
if err != nil {
t.Fatal(err)
}
if !bytes.Equal(raw, id) {
t.Errorf("Mask/UnMask didn't produce same result, got %x, expected %x", id, raw)
}
}
// -----------------------------------------------------------------------------
func TestMask8(t *testing.T) {
idm := Random8(crand.Reader, []byte("829bph7qn7KXc27NgQN1tAWuinYo0C2Q97bVk9sDN5SilqxHTgcBTMxQraMdbTD"))
fontain := sonyflake.NewSonyflake(sonyflake.Settings{})
uid, _ := fontain.NextID()
var raw [8]byte
binary.LittleEndian.PutUint64(raw[:], uid)
out, err := idm.Mask(raw[:])
if err != nil {
panic(err)
}
fmt.Printf("%d => %s\n", uid, base64.RawURLEncoding.EncodeToString(out))
id, err := idm.UnMask(out)
if err != nil {
panic(err)
}
if !bytes.Equal(raw[:], id) {
t.Errorf("Mask/UnMask didn't produce same result, got %x, expected %x", id, raw)
}
}
func BenchmarkMask8(b *testing.B) {
idm := Random8(crand.Reader, []byte("829bph7qn7KXc27NgQN1tAWuinYo0C2Q97bVk9sDN5SilqxHTgcBTMxQraMdbTD"))
fontain := sonyflake.NewSonyflake(sonyflake.Settings{})
uid, _ := fontain.NextID()
var raw [8]byte
binary.LittleEndian.PutUint64(raw[:], uid)
b.ReportAllocs()
for n := 0; n < b.N; n++ {
b.SetBytes(int64(len(raw)))
idm.Mask(raw[:])
}
}
func BenchmarkUnMask8(b *testing.B) {
idm := Random8(crand.Reader, []byte("829bph7qn7KXc27NgQN1tAWuinYo0C2Q97bVk9sDN5SilqxHTgcBTMxQraMdbTD"))
fontain := sonyflake.NewSonyflake(sonyflake.Settings{})
uid, _ := fontain.NextID()
var raw [8]byte
binary.LittleEndian.PutUint64(raw[:], uid)
out, _ := idm.Mask(raw[:])
b.ReportAllocs()
for n := 0; n < b.N; n++ {
b.SetBytes(int64(len(out)))
idm.UnMask(out)
}
}
Raw
output.txt
❯ go test -v
=== RUN Test16
=== RUN Test16/ClassicDeterministic
907f8535-ffba-480f-9957-19705a315cf1 => bK3as7RhS2glIKZK7oCqZ24e6blXUPJFGVB0Nws7NLU
=== RUN Test16/ClassicRandomized
6bc3c694-81fc-44ea-9448-2bd5178ea82d => 7m91MBDGAC3Zq70LNQLWPc2vKO8k83wGQ5By-45u__-EoXUpubwPjYkq1zGVC59x
=== RUN Test16/ModernDetermistic
fa0d2aa7-849b-48fb-98a4-871a83d803fa => qjNPbrXBigUoqu9ubnXYfUJCqvtjG4F9g7V1HYmpcbQ
=== RUN Test16/ModernRandomized
4f288956-6f8c-4249-88ff-b6f15cdc3e24 => hM-PlS3AT-XHg2APVieCd6aQa1LVFUeCaFfigLekokR6cEeCH1wN51yVBmLeGc6U
--- PASS: Test16 (0.00s)
--- PASS: Test16/ClassicDeterministic (0.00s)
--- PASS: Test16/ClassicRandomized (0.00s)
--- PASS: Test16/ModernDetermistic (0.00s)
--- PASS: Test16/ModernRandomized (0.00s)
=== RUN TestMask8
429435385135759654 => clCpCspxetjA_qUZOG_vmxHdRQnoVZM7
--- PASS: TestMask8 (0.00s)
PASS
ok idmask 0.218s
❯ go test -run=xx -bench=.
goos: darwin
goarch: arm64
pkg: idmask
Benchmark16/ClassicDeterministic/Mask-10 694686 1692 ns/op 9.46 MB/s 2600 B/op 36 allocs/op
Benchmark16/ClassicDeterministic/UnMask-10 699400 1679 ns/op 19.05 MB/s 2568 B/op 35 allocs/op
Benchmark16/ClassicRandomized/Mask-10 570679 2032 ns/op 7.87 MB/s 2616 B/op 36 allocs/op
Benchmark16/ClassicRandomized/UnMask-10 685022 1759 ns/op 27.29 MB/s 2568 B/op 35 allocs/op
Benchmark16/ModernDetermistic/Mask-10 1000000 1155 ns/op 13.85 MB/s 952 B/op 9 allocs/op
Benchmark16/ModernDetermistic/UnMask-10 1000000 1140 ns/op 28.08 MB/s 920 B/op 8 allocs/op
Benchmark16/ModernRandomized/Mask-10 790824 1464 ns/op 10.93 MB/s 968 B/op 9 allocs/op
Benchmark16/ModernRandomized/UnMask-10 1000000 1153 ns/op 41.62 MB/s 920 B/op 8 allocs/op
BenchmarkMask8-10 2454697 493.7 ns/op 16.20 MB/s 512 B/op 8 allocs/op
BenchmarkUnMask8-10 6952311 171.9 ns/op 139.60 MB/s 464 B/op 5 allocs/op
PASS
ok idmask 13.558s
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment