From 2ab4b4d53643197017715f4f6171c61dd8a2a9c8 Mon Sep 17 00:00:00 2001
From: Simon Ferquel <simon.ferquel@docker.com>
Date: Wed, 10 Jun 2020 15:07:23 +0200
Subject: [PATCH] Don't loose additional metadata fields

Signed-off-by: Simon Ferquel <simon.ferquel@docker.com>
---
 cli/command/context.go      | 45 +++++++++++++++++++++++++++++++++++--
 cli/command/context_test.go | 27 ++++++++++++++++++++++
 2 files changed, 70 insertions(+), 2 deletions(-)
 create mode 100644 cli/command/context_test.go

diff --git a/cli/command/context.go b/cli/command/context.go
index 0de75dc9a5..9c357a99d2 100644
--- a/cli/command/context.go
+++ b/cli/command/context.go
@@ -1,6 +1,7 @@
 package command
 
 import (
+	"encoding/json"
 	"errors"
 
 	"github.com/docker/cli/cli/context/store"
@@ -8,8 +9,48 @@ import (
 
 // DockerContext is a typed representation of what we put in Context metadata
 type DockerContext struct {
-	Description       string       `json:",omitempty"`
-	StackOrchestrator Orchestrator `json:",omitempty"`
+	Description       string
+	StackOrchestrator Orchestrator
+	AdditionalFields  map[string]interface{}
+}
+
+// MarshalJSON implements custom JSON marshalling
+func (dc DockerContext) MarshalJSON() ([]byte, error) {
+	s := map[string]interface{}{}
+	if dc.Description != "" {
+		s["Description"] = dc.Description
+	}
+	if dc.StackOrchestrator != "" {
+		s["StackOrchestrator"] = dc.StackOrchestrator
+	}
+	if dc.AdditionalFields != nil {
+		for k, v := range dc.AdditionalFields {
+			s[k] = v
+		}
+	}
+	return json.Marshal(s)
+}
+
+// UnmarshalJSON implements custom JSON marshalling
+func (dc *DockerContext) UnmarshalJSON(payload []byte) error {
+	var data map[string]interface{}
+	if err := json.Unmarshal(payload, &data); err != nil {
+		return err
+	}
+	for k, v := range data {
+		switch k {
+		case "Description":
+			dc.Description = v.(string)
+		case "StackOrchestrator":
+			dc.StackOrchestrator = Orchestrator(v.(string))
+		default:
+			if dc.AdditionalFields == nil {
+				dc.AdditionalFields = make(map[string]interface{})
+			}
+			dc.AdditionalFields[k] = v
+		}
+	}
+	return nil
 }
 
 // GetDockerContext extracts metadata from stored context metadata
diff --git a/cli/command/context_test.go b/cli/command/context_test.go
new file mode 100644
index 0000000000..28a33463be
--- /dev/null
+++ b/cli/command/context_test.go
@@ -0,0 +1,27 @@
+package command
+
+import (
+	"encoding/json"
+	"testing"
+
+	"gotest.tools/v3/assert"
+)
+
+func TestDockerContextMetadataKeepAdditionalFields(t *testing.T) {
+	c := DockerContext{
+		Description:       "test",
+		StackOrchestrator: OrchestratorSwarm,
+		AdditionalFields: map[string]interface{}{
+			"foo": "bar",
+		},
+	}
+	jsonBytes, err := json.Marshal(c)
+	assert.NilError(t, err)
+	assert.Equal(t, `{"Description":"test","StackOrchestrator":"swarm","foo":"bar"}`, string(jsonBytes))
+
+	var c2 DockerContext
+	assert.NilError(t, json.Unmarshal(jsonBytes, &c2))
+	assert.Equal(t, c2.AdditionalFields["foo"], "bar")
+	assert.Equal(t, c2.StackOrchestrator, OrchestratorSwarm)
+	assert.Equal(t, c2.Description, "test")
+}