refactor: 统一连接器的数据传送行为

Signed-off-by: mumuhhh <[email protected]>

Author: mumuhhh

pkg/connector/connector.go

@@ -18,12 +18,17 @@ package connector
 
 import (
 	"context"
+	"fmt"
 	"io"
 	"io/fs"
 	"net"
 	"os"
+	"path"
+	"path/filepath"
+	"strings"
 
 	"github.com/cockroachdb/errors"
+	"k8s.io/apimachinery/pkg/util/rand"
 	"k8s.io/klog/v2"
 
 	_const "github.com/kubesphere/kubekey/v4/pkg/const"
@@ -134,3 +139,23 @@ func isLocalIP(ipAddr string) bool {
 
 	return false
 }
+
+func PutData(ctx context.Context, data []byte, dest string, mode fs.FileMode, conn Connector) error {
+	dest = filepath.ToSlash(dest)
+	tmpDest := path.Join("/tmp", ".kk."+rand.String(10))
+
+	if err := conn.PutFile(ctx, data, tmpDest, mode); err != nil {
+		return err
+	}
+
+	// Escape single quotes to prevent shell injection
+	esc := func(s string) string {
+		return strings.ReplaceAll(s, "'", "'\\''")
+	}
+
+	cmd := fmt.Sprintf("mkdir -p '%s' && mv '%s' '%s' || { rm -f '%s'; exit 1; }",
+		esc(path.Dir(dest)), esc(tmpDest), esc(dest), esc(tmpDest))
+	_, _, err := conn.ExecuteCommand(ctx, cmd)
+
+	return err
+}

pkg/modules/copy/copy.go

@@ -28,7 +28,6 @@ import (
 	"github.com/cockroachdb/errors"
 	kkcorev1alpha1 "github.com/kubesphere/kubekey/api/core/v1alpha1"
 	"k8s.io/apimachinery/pkg/runtime"
-	"k8s.io/apimachinery/pkg/util/rand"
 	"k8s.io/klog/v2"
 	"k8s.io/utils/ptr"
 
@@ -251,15 +250,7 @@ func (ca copyArgs) copyAbsoluteDir(ctx context.Context, conn connector.Connector
 		}
 		dest := filepath.Join(ca.dest, rel)
 
-		tmpDest := filepath.Join("/tmp", ".kk."+rand.String(10))
-
-		if err = conn.PutFile(ctx, data, tmpDest, mode); err != nil {
-			return err
-		}
-
-		_, _, err = conn.ExecuteCommand(ctx, fmt.Sprintf("mkdir -p %s\nmv %s %s", filepath.Dir(dest), tmpDest, dest))
-
-		return err
+		return connector.PutData(ctx, data, dest, mode, conn)
 	})
 }
 
@@ -295,16 +286,7 @@ func (ca copyArgs) copyRelativeDir(ctx context.Context, pj project.Project, relP
 		}
 		dest := filepath.Join(ca.dest, rel)
 
-		tmpDest := filepath.Join("/tmp", ".kk."+rand.String(10))
-
-		err = conn.PutFile(ctx, data, tmpDest, mode)
-		if err != nil {
-			return err
-		}
-
-		_, _, err = conn.ExecuteCommand(ctx, fmt.Sprintf("mkdir -p %s\nmv %s %s", filepath.Dir(dest), tmpDest, dest))
-
-		return err
+		return connector.PutData(ctx, data, dest, mode, conn)
 	})
 }
 
@@ -330,14 +312,7 @@ func (ca copyArgs) copyContent(ctx context.Context, mode fs.FileMode, conn conne
 		mode = os.FileMode(*ca.mode)
 	}
 
-	tmpDest := filepath.Join("/tmp", ".kk."+rand.String(10))
-
-	if err := conn.PutFile(ctx, []byte(ca.content), tmpDest, mode); err != nil {
-		return internal.StdoutFailed, "failed to copy file", err
-	}
-
-	_, _, err := conn.ExecuteCommand(ctx, fmt.Sprintf("mkdir -p %s\nmv %s %s", filepath.Dir(ca.dest), tmpDest, ca.dest))
-
+	err := connector.PutData(ctx, []byte(ca.content), ca.dest, mode, conn)
 	if err != nil {
 		return internal.StdoutFailed, "failed to copy file", err
 	}
@@ -356,13 +331,6 @@ func (ca copyArgs) copyFile(ctx context.Context, data []byte, mode fs.FileMode,
 	if ca.mode != nil {
 		mode = os.FileMode(*ca.mode)
 	}
-	tmpDest := filepath.Join("/tmp", ".kk."+rand.String(10))
-
-	if err := conn.PutFile(ctx, data, tmpDest, mode); err != nil {
-		return err
-	}
-
-	_, _, err := conn.ExecuteCommand(ctx, fmt.Sprintf("mkdir -p %s\nmv %s %s", filepath.Dir(dest), tmpDest, dest))
 
-	return err
+	return connector.PutData(ctx, data, dest, mode, conn)
 }

pkg/modules/template/template.go

@@ -18,7 +18,6 @@ package template
 
 import (
 	"context"
-	"fmt"
 	"io/fs"
 	"math"
 	"os"
@@ -28,7 +27,6 @@ import (
 	"github.com/cockroachdb/errors"
 	kkcorev1alpha1 "github.com/kubesphere/kubekey/api/core/v1alpha1"
 	"k8s.io/apimachinery/pkg/runtime"
-	"k8s.io/apimachinery/pkg/util/rand"
 	"k8s.io/klog/v2"
 	"k8s.io/utils/ptr"
 
@@ -228,15 +226,8 @@ func handleRelativeDir(ctx context.Context, pj project.Project, relPath string,
 			}
 			dest = filepath.Join(ta.dest, rel)
 		}
-		tmpDest := filepath.Join("/tmp", ".kk."+rand.String(10))
 
-		if err = conn.PutFile(ctx, result, tmpDest, mode); err != nil {
-			return err
-		}
-
-		_, _, err = conn.ExecuteCommand(ctx, fmt.Sprintf("mkdir -p %s\nmv %s %s", filepath.Dir(dest), tmpDest, dest))
-
-		return err
+		return connector.PutData(ctx, result, dest, mode, conn)
 	})
 }
 
@@ -255,15 +246,8 @@ func (ta templateArgs) readFile(ctx context.Context, data string, mode fs.FileMo
 	if ta.mode != nil {
 		mode = os.FileMode(*ta.mode)
 	}
-	tmpDest := filepath.Join("/tmp", ".kk."+rand.String(10))
-
-	if err = conn.PutFile(ctx, result, tmpDest, mode); err != nil {
-		return err
-	}
 
-	_, _, err = conn.ExecuteCommand(ctx, fmt.Sprintf("mkdir -p %s\nmv %s %s", filepath.Dir(dest), tmpDest, dest))
-
-	return err
+	return connector.PutData(ctx, result, dest, mode, conn)
 }
 
 // absDir when template.src is absolute dir, get all files by os, parse it, and copy to remote.
@@ -304,15 +288,7 @@ func (ta templateArgs) absDir(ctx context.Context, conn connector.Connector, var
 			dest = filepath.Join(ta.dest, rel)
 		}
 
-		tmpDest := filepath.Join("/tmp", ".kk."+rand.String(10))
-
-		if err = conn.PutFile(ctx, result, tmpDest, mode); err != nil {
-			return err
-		}
-
-		_, _, err = conn.ExecuteCommand(ctx, fmt.Sprintf("mkdir -p %s\nmv %s %s", filepath.Dir(dest), tmpDest, dest))
-
-		return err
+		return connector.PutData(ctx, result, dest, mode, conn)
 	}); err != nil {
 		return errors.Wrapf(err, "failed to walk dir %q", ta.src)
 	}