diff --git a/util.go b/util.go index 8dc147a..8f3d8b6 100644 --- a/util.go +++ b/util.go @@ -42,7 +42,7 @@ import ( // <56-bytes-query> 0x80 (0x00 * 199) func pad(packet []byte) []byte { // get closest divisible by 64 to + 1 byte for 0x80 - minQuestionSize := (len(packet)+1+63)/64 + 64 + minQuestionSize := len(packet) + 1 + (64 - (len(packet)+1)%64) // padded size can't be less than minUDPQuestionSize minQuestionSize = max(minUDPQuestionSize, minQuestionSize) diff --git a/util_test.go b/util_test.go new file mode 100644 index 0000000..6c33263 --- /dev/null +++ b/util_test.go @@ -0,0 +1,31 @@ +package dnscrypt + +import ( + "crypto/rand" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestPadUnpad(t *testing.T) { + longBuf := make([]byte, 272) + _, err := rand.Read(longBuf) + require.NoError(t, err) + + tests := []struct { + packet []byte + expPaddedLen int + }{ + {[]byte("Example Test DNS packet"), 256}, + {longBuf, 320}, + } + for i, test := range tests { + padded := pad(test.packet) + assert.Equal(t, test.expPaddedLen, len(padded), "test %d", i) + + unpadded, err := unpad(padded) + assert.Nil(t, err, "test %d", i) + assert.Equal(t, test.packet, unpadded, "test %d", i) + } +}