sonic-net
diff --git a/‎gnmi_server/gnoi_file.go‎
Lines changed: 7 additions & 4 deletions b/‎gnmi_server/gnoi_file.go‎
Lines changed: 7 additions & 4 deletions
diff --git a/‎gnmi_server/gnoi_file_test.go‎
Lines changed: 144 additions & 20 deletions b/‎gnmi_server/gnoi_file_test.go‎
Lines changed: 144 additions & 20 deletions
diff --git a/‎gnmi_server/gnoi_system.go‎
Lines changed: 21 additions & 82 deletions b/‎gnmi_server/gnoi_system.go‎
Lines changed: 21 additions & 82 deletions
@@ -97,27 +97,30 @@ func (srv *FileServer) Get(req *gnoi_file_pb.GetRequest, stream gnoi_file_pb.Fil
 	return status.Errorf(codes.Unimplemented, "Method file.Get is unimplemented.")
 }
 
-// TransferToRemote RPC is unimplemented.
+// TransferToRemote downloads a file from a remote URL.
+// If DPU headers are present (HandleOnNPU mode), it downloads to NPU then uploads to the specified DPU.
 func (srv *FileServer) TransferToRemote(ctx context.Context, req *gnoi_file_pb.TransferToRemoteRequest) (*gnoi_file_pb.TransferToRemoteResponse, error) {
 	log.Infof("GNOI File TransferToRemote RPC called with request: %+v", req)
 	_, err := authenticate(srv.config, ctx, "gnoi", false)
 	if err != nil {
 		log.Errorf("authentication failed in TransferToRemote RPC: %v", err)
 		return nil, err
 	}
+
+	// Delegate all logic to the pure handler function
 	return gnoifile.HandleTransferToRemote(ctx, req)
 }
 
-// Put RPC is unimplemented.
+// Put implements the gNOI File.Put RPC.
+// It authenticates the request and delegates to the pure Go handler.
 func (srv *FileServer) Put(stream gnoi_file_pb.File_PutServer) error {
 	log.Infof("GNOI File Put RPC called")
 	_, err := authenticate(srv.config, stream.Context(), "gnoi", false)
 	if err != nil {
 		log.Errorf("authentication failed in Put RPC: %v", err)
 		return err
 	}
-	log.Warning("file.Put RPC is unimplemented")
-	return status.Errorf(codes.Unimplemented, "Method file.Put is unimplemented.")
+	return gnoifile.HandlePut(stream)
 }
 
 // Remove implements the corresponding RPC.
 
@@ -10,18 +10,28 @@ import (
 	"github.com/agiledragon/gomonkey/v2"
 	gnoi_common "github.com/openconfig/gnoi/common"
 	gnoi_file_pb "github.com/openconfig/gnoi/file"
+	gnoifile "github.com/sonic-net/sonic-gnmi/pkg/gnoi/file"
 	ssc "github.com/sonic-net/sonic-gnmi/sonic_service_client"
 
 	"github.com/stretchr/testify/assert"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/status"
 )
 
 // === Test Setup Helpers ===
-func createFileServer(t *testing.T, port int) *grpc.Server {
-	listener, err := net.Listen("tcp", fmt.Sprintf(":%d", port))
+func createFileServer(t *testing.T, port int) (*grpc.Server, string) {
+	var listener net.Listener
+	var err error
+
+	if port == 0 {
+		// Use dynamic port
+		listener, err = net.Listen("tcp", ":0")
+	} else {
+		listener, err = net.Listen("tcp", fmt.Sprintf(":%d", port))
+	}
 	if err != nil {
 		t.Fatalf("Failed to listen: %v", err)
 	}
@@ -40,19 +50,19 @@ func createFileServer(t *testing.T, port int) *grpc.Server {
 		}
 	}()
 
-	return s
+	return s, listener.Addr().String()
 }
 
 // === Actual Tests ===
 func TestGnoiFileServer(t *testing.T) {
-	s := createFileServer(t, 8081)
+	s, addr := createFileServer(t, 0) // Use dynamic port
 	defer s.Stop()
 
 	//tlsConfig := &tls.Config{InsecureSkipVerify: true}
 	//opts := []grpc.DialOption{grpc.WithTransportCredentials(credentials.NewTLS(tlsConfig))}
 	opts := []grpc.DialOption{grpc.WithTransportCredentials(insecure.NewCredentials())}
 
-	conn, err := grpc.Dial("127.0.0.1:8081", opts...)
+	conn, err := grpc.Dial(addr, opts...)
 	if err != nil {
 		t.Fatalf("Failed to dial server: %v", err)
 	}
@@ -188,21 +198,6 @@ func TestGnoiFileServer(t *testing.T) {
 		assert.Contains(t, err.Error(), "invalid syntax")
 	})
 
-	t.Run("Put Fails with Unimplemented Error", func(t *testing.T) {
-		patch := gomonkey.ApplyFuncReturn(authenticate, nil, nil)
-		defer patch.Reset()
-
-		putStream, err := client.Put(context.Background())
-		if err != nil {
-			t.Fatalf("Failed to create Put stream: %v", err)
-		}
-
-		// Expect Unimplemented error on CloseAndRecv
-		_, err = putStream.CloseAndRecv()
-		if err == nil || status.Code(err) != codes.Unimplemented {
-			t.Fatalf("Expected Unimplemented error, got: %v", err)
-		}
-	})
 	t.Run("Put Fails with Auth Error", func(t *testing.T) {
 		patch := gomonkey.ApplyFuncReturn(authenticate, nil, status.Error(codes.Unauthenticated, "unauthenticated"))
 		defer patch.Reset()
@@ -236,6 +231,71 @@ func TestGnoiFileServer(t *testing.T) {
 		}
 	})
 
+	t.Run("TransferToRemote DPU Success", func(t *testing.T) {
+		patches := gomonkey.NewPatches()
+		defer patches.Reset()
+
+		// Mock authenticate to succeed
+		patches.ApplyFuncReturn(authenticate, nil, nil)
+
+		// Mock HandleTransferToRemoteForDPUStreaming to succeed
+		patches.ApplyFunc(gnoifile.HandleTransferToRemoteForDPUStreaming,
+			func(ctx context.Context, req *gnoi_file_pb.TransferToRemoteRequest, dpuIndex string, dpuAddr string) (*gnoi_file_pb.TransferToRemoteResponse, error) {
+				return &gnoi_file_pb.TransferToRemoteResponse{}, nil
+			})
+
+		fs := &FileServer{
+			Server: &Server{
+				config: &Config{},
+			},
+		}
+
+		// Create context with DPU metadata (lines 117, 120, 125-126)
+		md := metadata.New(map[string]string{
+			"x-sonic-ss-target-type":  "dpu",
+			"x-sonic-ss-target-index": "0",
+		})
+		ctx := metadata.NewIncomingContext(context.Background(), md)
+
+		req := &gnoi_file_pb.TransferToRemoteRequest{
+			LocalPath: "/tmp/test.txt",
+		}
+
+		resp, err := fs.TransferToRemote(ctx, req)
+		assert.NoError(t, err)
+		assert.NotNil(t, resp)
+	})
+
+	t.Run("TransferToRemote NPU Success", func(t *testing.T) {
+		patches := gomonkey.NewPatches()
+		defer patches.Reset()
+
+		// Mock authenticate to succeed
+		patches.ApplyFuncReturn(authenticate, nil, nil)
+
+		// Mock HandleTransferToRemote to succeed
+		patches.ApplyFunc(gnoifile.HandleTransferToRemote,
+			func(ctx context.Context, req *gnoi_file_pb.TransferToRemoteRequest) (*gnoi_file_pb.TransferToRemoteResponse, error) {
+				return &gnoi_file_pb.TransferToRemoteResponse{}, nil
+			})
+
+		fs := &FileServer{
+			Server: &Server{
+				config: &Config{},
+			},
+		}
+
+		ctx := context.Background() // No DPU metadata - should call regular function
+
+		req := &gnoi_file_pb.TransferToRemoteRequest{
+			LocalPath: "/tmp/test.txt",
+		}
+
+		resp, err := fs.TransferToRemote(ctx, req)
+		assert.NoError(t, err)
+		assert.NotNil(t, resp)
+	})
+
 	t.Run("TransferToRemote Fails with Auth Error", func(t *testing.T) {
 		patch := gomonkey.ApplyFuncReturn(authenticate, nil, status.Error(codes.Unauthenticated, "unauthenticated"))
 		defer patch.Reset()
@@ -399,4 +459,68 @@ func TestGnoiFileServer(t *testing.T) {
 		}
 	})
 
+	// Test Put function success path (line 143)
+	t.Run("Put_Success", func(t *testing.T) {
+		patches := gomonkey.NewPatches()
+		defer patches.Reset()
+
+		patches.ApplyFuncReturn(authenticate, nil, nil)
+
+		// Mock the gnoifile.HandlePut function to return success
+		patches.ApplyFunc(gnoifile.HandlePut,
+			func(stream gnoi_file_pb.File_PutServer) error {
+				return nil
+			})
+
+		fs := &FileServer{
+			Server: &Server{
+				config: &Config{},
+			},
+		}
+
+		// Create a mock stream
+		mockStream := &mockPutStream{
+			ctx: context.Background(),
+		}
+
+		err := fs.Put(mockStream)
+		assert.NoError(t, err)
+	})
+
+}
+
+// Mock stream for Put testing
+type mockPutStream struct {
+	ctx context.Context
+}
+
+func (m *mockPutStream) Context() context.Context {
+	return m.ctx
+}
+
+func (m *mockPutStream) SendMsg(msg interface{}) error {
+	return nil
+}
+
+func (m *mockPutStream) RecvMsg(msg interface{}) error {
+	return nil
+}
+
+func (m *mockPutStream) Recv() (*gnoi_file_pb.PutRequest, error) {
+	return nil, nil
+}
+
+func (m *mockPutStream) SendAndClose(resp *gnoi_file_pb.PutResponse) error {
+	return nil
+}
+
+func (m *mockPutStream) SendHeader(metadata.MD) error {
+	return nil
+}
+
+func (m *mockPutStream) SetTrailer(metadata.MD) {
+}
+
+func (m *mockPutStream) SetHeader(metadata.MD) error {
+	return nil
 }
@@ -12,6 +12,7 @@ import (
 	"github.com/golang/protobuf/proto"
 	syspb "github.com/openconfig/gnoi/system"
 	"github.com/sonic-net/sonic-gnmi/common_utils"
+	"github.com/sonic-net/sonic-gnmi/pkg/gnoi/system"
 	ssc "github.com/sonic-net/sonic-gnmi/sonic_service_client"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
@@ -198,7 +199,20 @@ func (srv *Server) Reboot(ctx context.Context, req *syspb.RebootRequest) (*syspb
 		return nil, status.Errorf(codes.InvalidArgument, err.Error())
 	}
 
-	// Initialize State DB.
+	// Try the pure handler first (it handles DPU routing internally)
+	resp, err := system.HandleReboot(ctx, req)
+	if err != nil {
+		// If it's not the "local fallback" error, return the actual error
+		if status.Code(err) != codes.Unimplemented {
+			return nil, err
+		}
+		// Otherwise fall through to local handling
+	} else {
+		// DPU case handled successfully by pure handler
+		return resp, nil
+	}
+
+	// Initialize State DB for local reboot.
 	rclient, err := common_utils.GetRedisDBClient()
 	if err != nil {
 		return nil, status.Errorf(codes.Internal, err.Error())
@@ -211,15 +225,15 @@ func (srv *Server) Reboot(ctx context.Context, req *syspb.RebootRequest) (*syspb
 	}
 
 	// System reboot.
-	resp, err, _ := sendRebootReqOnNotifCh(ctx, req, rclient, rebootKey)
+	respInterface, err, _ := sendRebootReqOnNotifCh(ctx, req, rclient, rebootKey)
 	if err != nil {
 		return nil, err
 	}
-	if resp == nil {
+	if respInterface == nil {
 		log.V(2).Info("Reboot request received empty response from Reboot Backend.")
-		resp = &syspb.RebootResponse{}
+		respInterface = &syspb.RebootResponse{}
 	}
-	return resp.(*syspb.RebootResponse), nil
+	return respInterface.(*syspb.RebootResponse), nil
 }
 
 // RebootStatus implements the corresponding RPC.
@@ -308,83 +322,8 @@ func (srv *Server) SetPackage(rs syspb.System_SetPackageServer) error {
 	}
 	log.V(1).Info("gNOI: SetPackage request received")
 
-	// Create D-Bus client
-	dbus, err := ssc.NewDbusClient()
-	if err != nil {
-		log.Errorf("Failed to create D-Bus client: %v", err)
-		return status.Errorf(codes.Internal, "failed to create D-Bus client: %v", err)
-	}
-	defer dbus.Close()
-
-	// Receive the package request
-	req, err := rs.Recv()
-	if err != nil {
-		log.Errorf("Failed to receive package request: %v", err)
-		return err
-	}
-
-	// Validate request type
-	pkg, ok := req.GetRequest().(*syspb.SetPackageRequest_Package)
-	if !ok {
-		errMsg := fmt.Sprintf("invalid request type: %T, expected SetPackageRequest_Package", req.GetRequest())
-		log.Errorf(errMsg)
-		return status.Errorf(codes.InvalidArgument, errMsg)
-	}
-
-	// A filename and a version must be provided
-	if pkg.Package.Filename == "" {
-		log.Errorf("Filename is missing in package request")
-		return status.Errorf(codes.InvalidArgument, "filename is missing in package request")
-	}
-	if pkg.Package.Version == "" {
-		log.Errorf("Version is missing in package request")
-		return status.Errorf(codes.InvalidArgument, "version is missing in package request")
-	}
-	// Log the package filename and version
-	log.V(1).Infof("Package filename: %s, version: %s", pkg.Package.Filename, pkg.Package.Version)
-
-	// Download the package if RemoteDownload is provided
-	if pkg.Package.RemoteDownload != nil {
-		// Validate RemoteDownload
-		log.V(1).Infof("RemoteDownload provided")
-		// Check if the path is provided
-		if pkg.Package.RemoteDownload.Path == "" {
-			log.Errorf("RemoteDownload path is missing")
-			return status.Errorf(codes.InvalidArgument, "remote download path is missing")
-		}
-		log.V(1).Infof("RemoteDownload path: %s", pkg.Package.RemoteDownload.Path)
-
-		// Download the package
-		err = dbus.DownloadImage(pkg.Package.RemoteDownload.Path, pkg.Package.Filename)
-		if err != nil {
-			log.Errorf("Failed to download image: %v", err)
-			return status.Errorf(codes.Internal, "failed to download image: %v", err)
-		}
-		log.V(1).Infof("Package %s downloaded successfully to %s", pkg.Package.Version, pkg.Package.Filename)
-	}
-
-	// If activate is requested, install the package and set it to be the next boot image
-	if pkg.Package.Activate {
-		log.V(1).Infof("Activate is requested")
-		// Install the package
-		err = dbus.InstallImage(pkg.Package.Filename)
-		if err != nil {
-			log.Errorf("Failed to install image: %v", err)
-			return status.Errorf(codes.Internal, "failed to install image: %v", err)
-		}
-		log.V(1).Infof("Package %s installed successfully", pkg.Package.Filename)
-		// Currently, Installing the image will automatically set it as the next boot image
-		log.V(1).Infof("Package %s set as next boot image", pkg.Package.Filename)
-	}
-
-	// Send response to client
-	if err := rs.SendAndClose(&syspb.SetPackageResponse{}); err != nil {
-		log.Errorf("Failed to send response: %v", err)
-		return err
-	}
-
-	log.V(1).Infof("SetPackage completed successfully for %s", pkg.Package.Filename)
-	return nil
+	// Delegate all logic to the pure handler
+	return system.HandleSetPackageStream(rs)
 }
 
 // SwitchControlProcessor implements the corresponding RPC.