From 20f89208664fa830d128aadc32c278198f2d6747 Mon Sep 17 00:00:00 2001 From: benduncan <5178243+benduncan@users.noreply.github.com> Date: Thu, 9 Oct 2025 14:07:23 +0800 Subject: [PATCH] Improved validation and unit-testing for framework --- hive/daemon/daemon.go | 16 +-- hive/gateway/ec2/instance/RunInstances.go | 76 ++++++++++---- .../gateway/ec2/instance/RunInstances_test.go | 99 ++++++++++++++++--- hive/utils/utils.go | 37 +++++++ 4 files changed, 191 insertions(+), 37 deletions(-) diff --git a/hive/daemon/daemon.go b/hive/daemon/daemon.go index c3e07ec64..f8f9de127 100644 --- a/hive/daemon/daemon.go +++ b/hive/daemon/daemon.go @@ -24,6 +24,7 @@ import ( "text/template" "time" + "github.com/aws/aws-sdk-go/service/ec2" "github.com/kdomanski/iso9660" "github.com/mulgadc/hive/hive/config" "github.com/mulgadc/hive/hive/qmp" @@ -632,19 +633,16 @@ func (d *Daemon) handleEC2Launch(msg *nats.Msg) { log.Printf("Received message on subject: %s", msg.Subject) log.Printf("Message data: %s", string(msg.Data)) + var RunInstancesInput ec2.RunInstancesInput + //var Reservation ec2.Reservation + var ec2Req EC2Request var ec2Response config.EC2Response - if err := json.Unmarshal(msg.Data, &ec2Req); err != nil { + if err := json.Unmarshal(msg.Data, &RunInstancesInput); err != nil { ec2Response.Error = fmt.Sprintf("Error unmarshaling EC2 request: %v", err) ec2Response.Respond(msg) - return - } - - // Only handle RunInstances action - if ec2Req.Action != "RunInstances" { - ec2Response.Error = fmt.Sprintf("Ignoring non-RunInstances action: %s", ec2Req.Action) - ec2Response.Respond(msg) + slog.Error("Request does not match RunInstancesInput", "err", err) return } @@ -655,6 +653,7 @@ func (d *Daemon) handleEC2Launch(msg *nats.Msg) { if !exists { ec2Response.Error = fmt.Sprintf("Unsupported instance type: %s", ec2Req.InstanceType) ec2Response.Respond(msg) + slog.Error("handleEC2Launch", "err") return } @@ -1614,6 +1613,7 @@ func (d *Daemon) StartInstance(instance *vm.VM) error { drive.File = v.NBDURI // Cleanup hostname to point to nbd://localhost from [::] + // TODO: Make NBD host config defined, or remote NBD server if not running locally. drive.File = strings.Replace(drive.File, "[::]", "nbd://127.0.0.1", 1) if v.Boot { diff --git a/hive/gateway/ec2/instance/RunInstances.go b/hive/gateway/ec2/instance/RunInstances.go index c78391dce..d2f0dc293 100644 --- a/hive/gateway/ec2/instance/RunInstances.go +++ b/hive/gateway/ec2/instance/RunInstances.go @@ -1,13 +1,16 @@ package gateway_ec2_instance import ( + "bytes" "encoding/json" "errors" "fmt" + "log/slog" "strings" "github.com/aws/aws-sdk-go/aws" "github.com/aws/aws-sdk-go/service/ec2" + "github.com/mulgadc/hive/hive/utils" ) type RunInstancesResponse struct { @@ -128,7 +131,7 @@ Sample JSON: }, */ -func RunInstances(input *ec2.RunInstancesInput) (reservation ec2.Reservation, err error) { +func ValidateRunInstancesInput(input *ec2.RunInstancesInput) (err error) { // Check required args from JSON above // required:[ @@ -136,32 +139,53 @@ func RunInstances(input *ec2.RunInstancesInput) (reservation ec2.Reservation, er // "MinCount" // ] + if input == nil { + return errors.New("MissingParameter") + } + + if input.MinCount == nil { + return errors.New("MissingParameter") + } + if *input.MinCount == 0 { - return reservation, errors.New("InvalidParameterValue") + return errors.New("InvalidParameterValue") } if *input.MaxCount == 0 { - return reservation, errors.New("InvalidParameterValue") + return errors.New("InvalidParameterValue") } // Additional validation from EC2 spec if *input.MinCount > *input.MaxCount { - return reservation, errors.New("InvalidParameterValue") + return errors.New("InvalidParameterValue") } - if *input.ImageId == "" { - return reservation, errors.New("MissingParameter") + if input.ImageId == nil || *input.ImageId == "" { + return errors.New("MissingParameter") } - if *input.InstanceType == "" { - return reservation, errors.New("MissingParameter") + if input.InstanceType == nil || *input.InstanceType == "" { + return errors.New("MissingParameter") } - if !strings.HasPrefix(*aws.String(*input.ImageId), "ami-") { - return reservation, errors.New("InvalidAMIID.Malformed") + if !strings.HasPrefix(*input.ImageId, "ami-") { + return errors.New("InvalidAMIID.Malformed") } + return + +} + +func RunInstances(input *ec2.RunInstancesInput) (reservation ec2.Reservation, err error) { + + // Validate input + err = ValidateRunInstancesInput(input) + + if err != nil { + return reservation, err + } + // Marshal to JSON, to send over the wire for processing (NATS) jsonData, err := json.Marshal(input) if err != nil { @@ -169,10 +193,15 @@ func RunInstances(input *ec2.RunInstancesInput) (reservation ec2.Reservation, er } // Run the simulated JSON request via NATS, which will return a JSON response - jsonResp, err := EC2_Process_RunInstances(jsonData) + jsonResp := EC2_Process_RunInstances(jsonData) + + // Validate if the response is successful or an error + + responseError, err := utils.ValidateErrorPayload(jsonResp) if err != nil { - return reservation, fmt.Errorf("failed to process RunInstances request: %w", err) + slog.Error("Runinstances failed", "err", responseError.Code) + return reservation, errors.New(*responseError.Code) } // Unmarshal the JSON response back into a Reservation struct @@ -185,12 +214,24 @@ func RunInstances(input *ec2.RunInstancesInput) (reservation ec2.Reservation, er } -func EC2_Process_RunInstances(jsonData []byte) (output []byte, err error) { +func EC2_Process_RunInstances(jsonData []byte) (output []byte) { var input ec2.RunInstancesInput - err = json.Unmarshal(jsonData, &input) + + decoder := json.NewDecoder(bytes.NewReader(jsonData)) + decoder.DisallowUnknownFields() + + err := decoder.Decode(&input) if err != nil { - return output, fmt.Errorf("failed to unmarshal JSON to RunInstancesInput: %w", err) + // TODO: Move error codes with vars to errors.go + return utils.GenerateErrorPayload("ValidationError") + } + + // Ensure the payload provided the fields that EC2 expects before proceeding. + err = ValidateRunInstancesInput(&input) + + if err != nil { + return utils.GenerateErrorPayload("ValidationError") } // Here you would add the logic to actually create the instance in your system. @@ -216,9 +257,10 @@ func EC2_Process_RunInstances(jsonData []byte) (output []byte, err error) { // Return as JSON, to simulate the NATS response jsonResponse, err := json.Marshal(reservation) if err != nil { - return output, fmt.Errorf("failed to marshal reservation to JSON: %w", err) + slog.Error("EC2_Process_RunInstances could not marshal reservation", "err", err) + return nil } - return jsonResponse, nil + return jsonResponse } diff --git a/hive/gateway/ec2/instance/RunInstances_test.go b/hive/gateway/ec2/instance/RunInstances_test.go index 4b2f9c1fd..282142964 100644 --- a/hive/gateway/ec2/instance/RunInstances_test.go +++ b/hive/gateway/ec2/instance/RunInstances_test.go @@ -1,28 +1,30 @@ package gateway_ec2_instance import ( + "encoding/json" "errors" "strings" "testing" "github.com/aws/aws-sdk-go/aws" "github.com/aws/aws-sdk-go/service/ec2" + "github.com/mulgadc/hive/hive/utils" "github.com/stretchr/testify/assert" ) -func TestParseRunInstances(t *testing.T) { +var defaults = ec2.RunInstancesInput{ + ImageId: aws.String("ami-0abcdef1234567890"), + InstanceType: aws.String("t2.micro"), + MinCount: aws.Int64(1), + MaxCount: aws.Int64(1), + KeyName: aws.String("my-key-pair"), + SecurityGroupIds: []*string{ + aws.String("sg-0123456789abcdef0"), + }, + SubnetId: aws.String("subnet-6e7f829e"), +} - var defaults = ec2.RunInstancesInput{ - ImageId: aws.String("ami-0abcdef1234567890"), - InstanceType: aws.String("t2.micro"), - MinCount: aws.Int64(1), - MaxCount: aws.Int64(1), - KeyName: aws.String("my-key-pair"), - SecurityGroupIds: []*string{ - aws.String("sg-0123456789abcdef0"), - }, - SubnetId: aws.String("subnet-6e7f829e"), - } +func TestParseRunInstances(t *testing.T) { // Group multiple tests tests := []struct { @@ -179,3 +181,76 @@ func TestParseRunInstances(t *testing.T) { // Additional test } + +func TestEC2ProcessRunInstances(t *testing.T) { + + var input ec2.DescribeInstancesInput + + input.DryRun = aws.Bool(true) + + // Marshal to JSON, to send over the wire for processing (NATS) + jsonData, err := json.Marshal(input) + + // Expect no marshal error, even for invalid payload + assert.NoError(t, err) + + // Run the simulated JSON request via NATS, which will return a JSON response + jsonResp := EC2_Process_RunInstances(jsonData) + + responseError, err := utils.ValidateErrorPayload(jsonResp) + + // Should successfully parse the generated error payload + assert.Error(t, err) + + // Expect correct error code + assert.Equal(t, "ValidationError", aws.StringValue(responseError.Code)) + + // Confirm the correct input struct, but default values incorrect. + emptyRunInstance := ec2.RunInstancesInput{ImageId: aws.String("ami-0abcdef1234567890")} + + jsonData, err = json.Marshal(emptyRunInstance) + + // Expect no marshal error + assert.NoError(t, err) + + // Run the simulated JSON request via NATS, which will return a JSON response + jsonResp = EC2_Process_RunInstances(jsonData) + + _, err = utils.ValidateErrorPayload(jsonResp) + + // Expect correct error code + assert.Equal(t, "ValidationError", aws.StringValue(responseError.Code)) + + // Should successfully parse the generated error payload + assert.Error(t, err) + + // Run expected "good" input + // Marshal to JSON, to send over the wire for processing (NATS) + jsonData, err = json.Marshal(defaults) + + // Expect no marshal error + assert.NoError(t, err) + + // Run the simulated JSON request via NATS, which will return a JSON response + jsonResp = EC2_Process_RunInstances(jsonData) + + var reservation ec2.Reservation + + _, err = utils.ValidateErrorPayload(jsonResp) + + // Should successfully parse the generated error payload + assert.NoError(t, err) + + // Unmarshal + err = json.Unmarshal(jsonResp, &reservation) + + assert.NoError(t, err) + + // Sanity check expected output matches + assert.Len(t, reservation.Instances, 1) + + if len(reservation.Instances) > 0 { + assert.Equal(t, defaults.InstanceType, reservation.Instances[0].InstanceType) + } + +} diff --git a/hive/utils/utils.go b/hive/utils/utils.go index 4abd3617c..c870683b7 100644 --- a/hive/utils/utils.go +++ b/hive/utils/utils.go @@ -2,6 +2,7 @@ package utils import ( "bytes" + "encoding/json" "encoding/xml" "errors" "fmt" @@ -13,7 +14,9 @@ import ( "syscall" "time" + "github.com/aws/aws-sdk-go/aws" "github.com/aws/aws-sdk-go/private/protocol/xml/xmlutil" + "github.com/aws/aws-sdk-go/service/ec2" ) func ReadPidFile(name string) (int, error) { @@ -235,3 +238,37 @@ func GenerateXMLPayload(locationName string, payload interface{}) interface{} { v.Field(0).Set(reflect.ValueOf(payload)) return v.Interface() } + +// Generate JSON Error Payload +func GenerateErrorPayload(code string) (jsonResponse []byte) { + + var responseError ec2.ResponseError + responseError.Code = aws.String(code) + + // Return as JSON, to simulate the NATS response + jsonResponse, err := json.Marshal(responseError) + if err != nil { + slog.Error("GenerateErrorPayload could not marshal JSON payload", "err", err) + return nil + } + + return + +} + +// Validate the payload is an ec2.ResponseError +func ValidateErrorPayload(payload []byte) (responseError ec2.ResponseError, err error) { + + decoder := json.NewDecoder(bytes.NewReader(payload)) + decoder.DisallowUnknownFields() + + err = decoder.Decode(&responseError) + + if err == nil { + // TODO: Move error codes with vars to errors.go + return responseError, errors.New("ResponseError detected") + } + + return responseError, nil + +}