Skip to content

Commit

Permalink
add support for gzipped metadata
Browse files Browse the repository at this point in the history
  • Loading branch information
eljohnson92 committed Sep 13, 2024
1 parent e269c00 commit 7f52b69
Show file tree
Hide file tree
Showing 4 changed files with 126 additions and 23 deletions.
1 change: 1 addition & 0 deletions go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ go 1.21

require (
github.com/go-resty/resty/v2 v2.14.0
github.com/jarcoal/httpmock v1.3.1
github.com/stretchr/testify v1.9.0
)

Expand Down
4 changes: 4 additions & 0 deletions go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,10 @@ github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSs
github.com/go-resty/resty/v2 v2.14.0 h1:/rhkzsAqGQkozwfKS5aFAbb6TyKd3zyFRWcdRXLPCAU=
github.com/go-resty/resty/v2 v2.14.0/go.mod h1:IW6mekUOsElt9C7oWr0XRt9BNSD6D5rr9mhk6NjmNHg=
github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
github.com/jarcoal/httpmock v1.3.1 h1:iUx3whfZWVf3jT01hQTO/Eo5sAYtB2/rqaUuOtpInww=
github.com/jarcoal/httpmock v1.3.1/go.mod h1:3yb8rc4BI7TCBhFY8ng0gjuLKJNquuDNiPaZjnENuYg=
github.com/maxatome/go-testdeep v1.12.0 h1:Ql7Go8Tg0C1D/uMMX59LAoYK7LffeJQ6X2T04nTH68g=
github.com/maxatome/go-testdeep v1.12.0/go.mod h1:lPZc/HAcJMP92l7yI6TRz1aZN5URwUBUAfUNvrclaNM=
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=
Expand Down
31 changes: 29 additions & 2 deletions userdata.go
Original file line number Diff line number Diff line change
@@ -1,13 +1,18 @@
package metadata

import (
"bytes"
"compress/gzip"
"context"
"encoding/base64"
"fmt"
"io"
)

var gzipMagic = []byte{0x1F, 0x8B, 0x08}

// GetUserData returns the user data for the current instance.
// NOTE: The result of this endpoint is automatically decoded from base64.
// NOTE: The result of this endpoint is automatically decoded from base64 and un-gzipped if needed.
func (c *Client) GetUserData(ctx context.Context) (string, error) {
req := c.R(ctx)

Expand All @@ -21,6 +26,28 @@ func (c *Client) GetUserData(ctx context.Context) (string, error) {
if err != nil {
return "", fmt.Errorf("failed to decode user-data: %w", err)
}
rawBytes, err := ungzipIfNeeded(decodedBytes)
if err != nil {
return "", fmt.Errorf("failed to ungzip user-data: %w", err)
}
return string(rawBytes), nil
}

return string(decodedBytes), nil
// hasGzipMagicNumber checks for the gzipMagic bytes at the beginning of the source
func hasGzipMagicNumber(source []byte) bool {
return bytes.HasPrefix(source, gzipMagic)
}

// ungzipIfNeeded checks for the gzip magic number and unzips the bytes if necessary,
// otherwise it returns the original raw bytes
func ungzipIfNeeded(raw []byte) ([]byte, error) {
if !hasGzipMagicNumber(raw) {
return raw, nil
}
reader, err := gzip.NewReader(bytes.NewReader(raw))
if err != nil {
return nil, err
}
defer reader.Close()
return io.ReadAll(reader)
}
113 changes: 92 additions & 21 deletions userdata_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,44 +3,115 @@ package metadata
import (
"context"
"encoding/base64"
"errors"
"fmt"
"net/http"
"testing"

"github.com/jarcoal/httpmock"

"github.com/stretchr/testify/assert"
)

type UserdataMockClient struct {
UserData string
GetUserDataError error
}
var mockMetadataHost = fmt.Sprintf("%s://%s/%s", APIProto, APIHost, APIVersion)

func (m *UserdataMockClient) GetUserData(ctx context.Context) (string, error) {
if m.GetUserDataError != nil {
return "", m.GetUserDataError
}
return m.UserData, nil
func SetupMockClient() *http.Client {
// create mock client
mockClient := http.DefaultClient
httpmock.ActivateNonDefault(mockClient)

// Mock out token request
tokenResponder := httpmock.NewStringResponder(200, "[\"token\"]")
httpmock.RegisterResponder("PUT", fmt.Sprintf("%s/token", mockMetadataHost), tokenResponder)
return mockClient
}

func TestGetUserData_Success(t *testing.T) {
mockClient := &UserdataMockClient{
UserData: base64.StdEncoding.EncodeToString([]byte("mock-user-data")),
}
mockClient := SetupMockClient()
// Mock out user-data response with the encoded value for "mock-user-data"
instanceResponder := httpmock.NewStringResponder(200, "bW9jay11c2VyLWRhdGE=")
httpmock.RegisterResponder("GET", fmt.Sprintf("%s/user-data", mockMetadataHost), instanceResponder)

newClient, err := NewClient(context.Background(), func(options *clientCreateConfig) {
options.HTTPClient = mockClient
})
assert.NoError(t, err, "Expected no error")

userData, err := newClient.GetUserData(context.Background())

assert.NoError(t, err, "Expected no error")

userData, err := mockClient.GetUserData(context.Background())
assert.Equal(t, "mock-user-data", userData, "Unexpected user data")
}

func TestGetUserDataGzip_Success(t *testing.T) {
mockClient := SetupMockClient()
// Mock out user-data response with the gzipped encoded value for "mock-user-data"
instanceResponder := httpmock.NewStringResponder(200, "H4sIAO0n32YAA8vNT87WLS1OLdJNSSxJBACRtuznDgAAAA==")
httpmock.RegisterResponder("GET", fmt.Sprintf("%s/user-data", mockMetadataHost), instanceResponder)

newClient, err := NewClient(context.Background(), func(options *clientCreateConfig) {
options.HTTPClient = mockClient
})
assert.NoError(t, err, "Expected no error")

userData, err := newClient.GetUserData(context.Background())

assert.NoError(t, err, "Expected no error")
// Note "bW9jay11c2VyLWRhdGE=" is the encoded value
assert.Equal(t, "bW9jay11c2VyLWRhdGE=", userData, "Unexpected user data")

assert.Equal(t, "mock-user-data", userData, "Unexpected user data")
}

func TestGetUserDataGzip_Error(t *testing.T) {
mockClient := SetupMockClient()
// Mock out user-data response with the invalid gzip encoded value for "mock-user-data"
invalidGzipData := []byte{0x1F, 0x8B, 0x08, 0x23}
userDataResponse := base64.StdEncoding.EncodeToString(invalidGzipData)
instanceResponder := httpmock.NewStringResponder(200, userDataResponse)
httpmock.RegisterResponder("GET", fmt.Sprintf("%s/user-data", mockMetadataHost), instanceResponder)

newClient, err := NewClient(context.Background(), func(options *clientCreateConfig) {
options.HTTPClient = mockClient
})
assert.NoError(t, err, "Expected no error")

userData, err := newClient.GetUserData(context.Background())

assert.EqualErrorf(t, err, "failed to ungzip user-data: unexpected EOF", "Unexpected error message")

assert.Equal(t, "", userData, "expected Empty Userdata")
}

func TestGetUserData_Error(t *testing.T) {
mockClient := &UserdataMockClient{
GetUserDataError: errors.New("mock error"),
}
mockClient := SetupMockClient()

instanceResponder := httpmock.NewStringResponder(500, "{\"errors\": [{\"reason\": \"failed to get metadata\"}]}")
httpmock.RegisterResponder("GET", fmt.Sprintf("%s/user-data", mockMetadataHost), instanceResponder)
newClient, err := NewClient(context.Background(), func(options *clientCreateConfig) {
options.HTTPClient = mockClient
})
assert.NoError(t, err, "Expected no error")

userData, err := newClient.GetUserData(context.Background())

assert.Error(t, err, "Expected an error")
assert.Equal(t, "", userData, "Expected empty user data")
assert.EqualErrorf(t, err, "[500] failed to get metadata", "Unexpected error message")
}

func TestGetUserDataDecode_Error(t *testing.T) {
mockClient := SetupMockClient()
// Mock out user-data response with the gzipped encoded value for "mock-user-data"
instanceResponder := httpmock.NewStringResponder(200, "invalid base64")
httpmock.RegisterResponder("GET", fmt.Sprintf("%s/user-data", mockMetadataHost), instanceResponder)

newClient, err := NewClient(context.Background(), func(options *clientCreateConfig) {
options.HTTPClient = mockClient
})
assert.NoError(t, err, "Expected no error")

userData, err := mockClient.GetUserData(context.Background())
userData, err := newClient.GetUserData(context.Background())

assert.Error(t, err, "Expected an error")
assert.Equal(t, "", userData, "Expected empty user data")
assert.EqualError(t, err, "mock error", "Unexpected error message")
assert.EqualErrorf(t, err, "failed to decode user-data: illegal base64 data at input byte 7", "Unexpected error message")
}

0 comments on commit 7f52b69

Please sign in to comment.