From 4c8baddefca5f59e4a614a700e6618bf64313817 Mon Sep 17 00:00:00 2001
From: Boshi Lian <farmer1992@gmail.com>
Date: Sun, 23 Oct 2022 08:07:30 +0000
Subject: [PATCH] fix cache set

---
 plugin/kubernetes/kubernetes.go | 8 ++++++--
 plugin/yaml/yaml.go             | 6 ++++--
 2 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/plugin/kubernetes/kubernetes.go b/plugin/kubernetes/kubernetes.go
index 8edc5b0d..591f4f12 100644
--- a/plugin/kubernetes/kubernetes.go
+++ b/plugin/kubernetes/kubernetes.go
@@ -7,6 +7,7 @@ import (
 	"fmt"
 	"net"
 	"regexp"
+	"time"
 
 	gocache "github.com/patrickmn/go-cache"
 	"github.com/tg123/sshpiper/libplugin"
@@ -69,6 +70,7 @@ func newKubernetesPlugin() (*plugin, error) {
 		k8sclient: k8sclient.CoreV1(),
 		lister:    lister,
 		stop:      stop,
+		cache:     gocache.New(1*time.Minute, 10*time.Minute),
 	}, nil
 }
 
@@ -154,7 +156,8 @@ func (p *plugin) createUpstream(conn libplugin.ConnMetadata, pipe *piperv1beta1.
 
 	if originPassword != "" {
 		u.Auth = libplugin.CreatePasswordAuth([]byte(originPassword))
-		return u, p.cache.Add(conn.UniqueID(), pipe, gocache.DefaultExpiration)
+		p.cache.Set(conn.UniqueID(), pipe, gocache.DefaultExpiration)
+		return u, nil
 	}
 
 	secret, err := p.k8sclient.Secrets(pipe.Namespace).Get(context.Background(), to.PrivateKeySecret.Name, metav1.GetOptions{})
@@ -164,7 +167,8 @@ func (p *plugin) createUpstream(conn libplugin.ConnMetadata, pipe *piperv1beta1.
 
 	data := secret.Data["privatekey"]
 	if data != nil {
-		return u, p.cache.Add(conn.UniqueID(), pipe, gocache.DefaultExpiration)
+		p.cache.Set(conn.UniqueID(), pipe, gocache.DefaultExpiration)
+		return u, nil
 	}
 
 	return nil, fmt.Errorf("no password or private key found")
diff --git a/plugin/yaml/yaml.go b/plugin/yaml/yaml.go
index 5bcf4fb8..2ba387c3 100644
--- a/plugin/yaml/yaml.go
+++ b/plugin/yaml/yaml.go
@@ -211,7 +211,8 @@ func (p *plugin) createUpstream(conn libplugin.ConnMetadata, to pipeConfigTo, or
 	// password found
 	if pass != "" {
 		u.Auth = libplugin.CreatePasswordAuth([]byte(pass))
-		return u, p.cache.Add(conn.UniqueID(), &to, cache.DefaultExpiration)
+		p.cache.Set(conn.UniqueID(), &to, cache.DefaultExpiration)
+		return u, nil
 	}
 
 	// try private key
@@ -225,7 +226,8 @@ func (p *plugin) createUpstream(conn libplugin.ConnMetadata, to pipeConfigTo, or
 
 	if data != nil {
 		u.Auth = libplugin.CreatePrivateKeyAuth(data)
-		return u, p.cache.Add(conn.UniqueID(), &to, cache.DefaultExpiration)
+		p.cache.Set(conn.UniqueID(), &to, cache.DefaultExpiration)
+		return u, nil
 	}
 
 	return nil, fmt.Errorf("no password or private key found")