From d03fd5805a56949992f6d9af5de5c69a9952cec5 Mon Sep 17 00:00:00 2001
From: Anthony Lapenna <lapenna.anthony@gmail.com>
Date: Fri, 23 Nov 2018 11:46:51 +1300
Subject: [PATCH] feat(api): support AGENT_SECRET environment variable (#2486)

---
 api/cmd/portainer/main.go                    |  3 ++-
 api/crypto/ecdsa.go                          | 23 +++++++++++++++++---
 api/docker/client.go                         |  2 +-
 api/exec/swarm_stack.go                      |  2 +-
 api/http/handler/websocket/websocket_exec.go |  3 ++-
 api/http/proxy/docker_transport.go           |  2 +-
 api/portainer.go                             |  2 +-
 7 files changed, 28 insertions(+), 9 deletions(-)

diff --git a/api/cmd/portainer/main.go b/api/cmd/portainer/main.go
index 0e0c189ac..f7890c2c4 100644
--- a/api/cmd/portainer/main.go
+++ b/api/cmd/portainer/main.go
@@ -2,6 +2,7 @@ package main // import "github.com/portainer/portainer"
 
 import (
 	"encoding/json"
+	"os"
 	"strings"
 	"time"
 
@@ -88,7 +89,7 @@ func initJWTService(authenticationEnabled bool) portainer.JWTService {
 }
 
 func initDigitalSignatureService() portainer.DigitalSignatureService {
-	return &crypto.ECDSAService{}
+	return crypto.NewECDSAService(os.Getenv("AGENT_SECRET"))
 }
 
 func initCryptoService() portainer.CryptoService {
diff --git a/api/crypto/ecdsa.go b/api/crypto/ecdsa.go
index 003547531..f2b87d639 100644
--- a/api/crypto/ecdsa.go
+++ b/api/crypto/ecdsa.go
@@ -8,6 +8,8 @@ import (
 	"encoding/base64"
 	"encoding/hex"
 	"math/big"
+
+	"github.com/portainer/portainer"
 )
 
 const (
@@ -26,6 +28,15 @@ type ECDSAService struct {
 	privateKey    *ecdsa.PrivateKey
 	publicKey     *ecdsa.PublicKey
 	encodedPubKey string
+	secret        string
+}
+
+// NewECDSAService returns a pointer to a ECDSAService.
+// An optional secret can be specified
+func NewECDSAService(secret string) *ECDSAService {
+	return &ECDSAService{
+		secret: secret,
+	}
 }
 
 // EncodedPublicKey returns the encoded version of the public that can be used
@@ -91,11 +102,17 @@ func (service *ECDSAService) GenerateKeyPair() ([]byte, []byte, error) {
 	return private, public, nil
 }
 
-// Sign creates a signature from a message.
-// It automatically hash the message using MD5 and creates a signature from
+// CreateSignature creates a digital signature.
+// It automatically hash a specific message using MD5 and creates a signature from
 // that hash.
 // It then encodes the generated signature in base64.
-func (service *ECDSAService) Sign(message string) (string, error) {
+func (service *ECDSAService) CreateSignature() (string, error) {
+
+	message := portainer.PortainerAgentSignatureMessage
+	if service.secret != "" {
+		message = service.secret
+	}
+
 	hash := HashFromBytes([]byte(message))
 
 	r := big.NewInt(0)
diff --git a/api/docker/client.go b/api/docker/client.go
index 9c608a19b..913e5dd87 100644
--- a/api/docker/client.go
+++ b/api/docker/client.go
@@ -67,7 +67,7 @@ func createAgentClient(endpoint *portainer.Endpoint, signatureService portainer.
 		return nil, err
 	}
 
-	signature, err := signatureService.Sign(portainer.PortainerAgentSignatureMessage)
+	signature, err := signatureService.CreateSignature()
 	if err != nil {
 		return nil, err
 	}
diff --git a/api/exec/swarm_stack.go b/api/exec/swarm_stack.go
index aa32bfe54..f41a581db 100644
--- a/api/exec/swarm_stack.go
+++ b/api/exec/swarm_stack.go
@@ -140,7 +140,7 @@ func (manager *SwarmStackManager) updateDockerCLIConfiguration(dataPath string)
 		return err
 	}
 
-	signature, err := manager.signatureService.Sign(portainer.PortainerAgentSignatureMessage)
+	signature, err := manager.signatureService.CreateSignature()
 	if err != nil {
 		return err
 	}
diff --git a/api/http/handler/websocket/websocket_exec.go b/api/http/handler/websocket/websocket_exec.go
index d797e020a..6f2d07d22 100644
--- a/api/http/handler/websocket/websocket_exec.go
+++ b/api/http/handler/websocket/websocket_exec.go
@@ -111,12 +111,13 @@ func (handler *Handler) proxyWebsocketRequest(w http.ResponseWriter, r *http.Req
 		}
 	}
 
-	signature, err := handler.SignatureService.Sign(portainer.PortainerAgentSignatureMessage)
+	signature, err := handler.SignatureService.CreateSignature()
 	if err != nil {
 		return err
 	}
 
 	proxy.Director = func(incoming *http.Request, out http.Header) {
+		out.Set(portainer.PortainerAgentPublicKeyHeader, handler.SignatureService.EncodedPublicKey())
 		out.Set(portainer.PortainerAgentSignatureHeader, signature)
 		out.Set(portainer.PortainerAgentTargetHeader, params.nodeName)
 	}
diff --git a/api/http/proxy/docker_transport.go b/api/http/proxy/docker_transport.go
index 37e46c2dd..2476685cc 100644
--- a/api/http/proxy/docker_transport.go
+++ b/api/http/proxy/docker_transport.go
@@ -64,7 +64,7 @@ func (p *proxyTransport) proxyDockerRequest(request *http.Request) (*http.Respon
 	request.URL.Path = path
 
 	if p.enableSignature {
-		signature, err := p.SignatureService.Sign(portainer.PortainerAgentSignatureMessage)
+		signature, err := p.SignatureService.CreateSignature()
 		if err != nil {
 			return nil, err
 		}
diff --git a/api/portainer.go b/api/portainer.go
index f75574723..ee2d6687f 100644
--- a/api/portainer.go
+++ b/api/portainer.go
@@ -626,7 +626,7 @@ type (
 		GenerateKeyPair() ([]byte, []byte, error)
 		EncodedPublicKey() string
 		PEMHeaders() (string, string)
-		Sign(message string) (string, error)
+		CreateSignature() (string, error)
 	}
 
 	// JWTService represents a service for managing JWT tokens