Addressing Review Comments and Channel Close Fix

Signed-off-by: Manish Ranjan Mahanta <mmahanta@microsoft.com>

Author: Manish Ranjan Mahanta

internal/gcs/bridge.go

@@ -332,7 +332,7 @@ func (brdg *bridge) recvLoop() error {
 			}
 
 		case prot.MsgTypeNotify:
-			if typ != prot.NotifyContainer|prot.MsgTypeNotify {
+			if typ != prot.NotifyContainer|prot.ComputeSystem|prot.MsgTypeNotify {
 				return fmt.Errorf("bridge received unknown unknown notification message %s", typ)
 			}
 			var ntf prot.ContainerNotification

internal/gcs/bridge_test.go

@@ -179,7 +179,7 @@ func notifyThroughBridge(t *testing.T, typ prot.MsgType, msg interface{}, fn not
 func TestBridgeNotify(t *testing.T) {
 	ntf := &prot.ContainerNotification{Operation: "testing"}
 	recvd := false
-	err := notifyThroughBridge(t, prot.MsgTypeNotify|prot.NotifyContainer, ntf, func(nntf *prot.ContainerNotification) error {
+	err := notifyThroughBridge(t, prot.MsgTypeNotify|prot.ComputeSystem|prot.NotifyContainer, ntf, func(nntf *prot.ContainerNotification) error {
 		if !reflect.DeepEqual(ntf, nntf) {
 			t.Errorf("%+v != %+v", ntf, nntf)
 		}
@@ -197,7 +197,7 @@ func TestBridgeNotify(t *testing.T) {
 func TestBridgeNotifyFailure(t *testing.T) {
 	ntf := &prot.ContainerNotification{Operation: "testing"}
 	errMsg := "notify should have failed"
-	err := notifyThroughBridge(t, prot.MsgTypeNotify|prot.NotifyContainer, ntf, func(nntf *prot.ContainerNotification) error {
+	err := notifyThroughBridge(t, prot.MsgTypeNotify|prot.ComputeSystem|prot.NotifyContainer, ntf, func(nntf *prot.ContainerNotification) error {
 		return errors.New(errMsg)
 	})
 	if err == nil || !strings.Contains(err.Error(), errMsg) {

internal/gcs/guestconnection.go

@@ -191,7 +191,7 @@ func (gc *GuestConnection) ModifyServiceSettings(ctx context.Context, serviceTyp
 
 	req := prot.ServiceModificationRequest{
 		RequestBase:  makeRequest(ctx, nullContainerID),
-		PropertyType: serviceType.String(),
+		PropertyType: string(serviceType),
 		Settings:     settings,
 	}
 	var resp prot.ResponseBase

internal/gcs/guestconnection_test.go

@@ -131,7 +131,7 @@ func simpleGcsLoop(t *testing.T, rw io.ReadWriter) error {
 				return err
 			}
 			time.Sleep(50 * time.Millisecond)
-			err = sendJSON(t, rw, prot.MsgType(prot.MsgTypeNotify|prot.NotifyContainer), 0, &prot.ContainerNotification{
+			err = sendJSON(t, rw, prot.MsgType(prot.MsgTypeNotify|prot.ComputeSystem|prot.NotifyContainer), 0, &prot.ContainerNotification{
 				RequestBase: prot.RequestBase{
 					ContainerID: req.ContainerID,
 				},

internal/gcs/prot/protocol.go

@@ -124,21 +124,12 @@ const (
 	RPCModifyServiceSettings RPCProc = ComputeService | (iota+1)<<8 | 1
 )
 
-type ServiceModifyPropertyType uint32
+type ServiceModifyPropertyType string
 
 const (
-	LogForwardService ServiceModifyPropertyType = iota
+	LogForwardService = ServiceModifyPropertyType("LogForwardService")
 )
 
-func (m ServiceModifyPropertyType) String() string {
-	switch m {
-	case LogForwardService:
-		return "LogForwardService"
-	default:
-		return fmt.Sprintf("UnknownModifyServiceType(%d)", m)
-	}
-}
-
 func (rpc RPCProc) String() string {
 	switch rpc {
 	case RPCCreate:

internal/uvm/create_wcow.go

@@ -266,11 +266,13 @@ func prepareCommonConfigDoc(ctx context.Context, uvm *UtilityVM, opts *OptionsWC
 	}
 
 	maps.Copy(doc.VirtualMachine.Devices.HvSocket.HvSocketConfig.ServiceTable, opts.AdditionalHyperVConfig)
-	key := prot.WindowsLoggingHvsockServiceID.String()
-	doc.VirtualMachine.Devices.HvSocket.HvSocketConfig.ServiceTable[key] = hcsschema.HvSocketServiceConfig{
-		AllowWildcardBinds:        true,
-		BindSecurityDescriptor:    "D:P(A;;FA;;;SY)(A;;FA;;;BA)",
-		ConnectSecurityDescriptor: "D:P(A;;FA;;;SY)(A;;FA;;;BA)",
+	if opts.ForwardLogs {
+		key := prot.WindowsLoggingHvsockServiceID.String()
+		doc.VirtualMachine.Devices.HvSocket.HvSocketConfig.ServiceTable[key] = hcsschema.HvSocketServiceConfig{
+			AllowWildcardBinds:        true,
+			BindSecurityDescriptor:    "D:P(A;;FA;;;SY)(A;;FA;;;BA)",
+			ConnectSecurityDescriptor: "D:P(A;;FA;;;SY)(A;;FA;;;BA)",
+		}
 	}
 
 	// Handle StorageQoS if set
@@ -551,14 +553,16 @@ func CreateWCOW(ctx context.Context, opts *OptionsWCOW) (_ *UtilityVM, err error
 		return nil, fmt.Errorf("error while creating the compute system: %w", err)
 	}
 
-	// Create a socket that the executed program can send to. This is usually
-	// used by Log Forward Service to send log data.
-	uvm.outputHandler = opts.OutputHandlerCreator(opts.Options)
-	uvm.outputProcessingDone = make(chan struct{})
-	uvm.outputListener, err = winio.ListenHvsock(&winio.HvsockAddr{
-		VMID:      uvm.RuntimeID(),
-		ServiceID: prot.WindowsLoggingHvsockServiceID,
-	})
+	if opts.ForwardLogs {
+		// Create a socket that the executed program can send to. This is usually
+		// used by Log Forward Service to send log data.
+		uvm.outputHandler = opts.OutputHandlerCreator(opts.Options)
+		uvm.outputProcessingDone = make(chan struct{})
+		uvm.outputListener, err = winio.ListenHvsock(&winio.HvsockAddr{
+			VMID:      uvm.RuntimeID(),
+			ServiceID: prot.WindowsLoggingHvsockServiceID,
+		})
+	}
 
 	gcsServiceID := prot.WindowsGcsHvsockServiceID
 	if opts.SecurityPolicyEnabled {

internal/uvm/log_wcow.go

@@ -12,57 +12,63 @@ import (
 )
 
 func (uvm *UtilityVM) StartLogForwarding(ctx context.Context) error {
-	// Implementation for stopping the log forwarder
-	if uvm.OS() == "windows" && uvm.gc != nil {
-		wcaps := gcs.GetWCOWCapabilities(uvm.gc.Capabilities())
-		if wcaps != nil && wcaps.IsLogForwardingSupported() {
-			req := guestrequest.LogForwardServiceRPCRequest{
-				RPCType:  guestrequest.RPCStartLogForwarding,
-				Settings: "",
-			}
-			err := uvm.gc.ModifyServiceSettings(ctx, prot.LogForwardService, req)
-			if err != nil {
-				return err
-			}
-		} else {
-			log.G(ctx).WithField("os", uvm.operatingSystem).Error("Log forwarding not supported for this OS")
+	// Implementation for starting the log forwarding service
+	if uvm.OS() != "windows" || uvm.gc == nil {
+		return errNotSupported
+	}
+
+	wcaps := gcs.GetWCOWCapabilities(uvm.gc.Capabilities())
+	if wcaps != nil && wcaps.IsLogForwardingSupported() {
+		req := guestrequest.LogForwardServiceRPCRequest{
+			RPCType:  guestrequest.RPCStartLogForwarding,
+			Settings: "",
+		}
+		err := uvm.gc.ModifyServiceSettings(ctx, prot.LogForwardService, req)
+		if err != nil {
+			return err
 		}
+	} else {
+		log.G(ctx).WithField("os", uvm.operatingSystem).Error("Log forwarding not supported for this OS")
 	}
 	return nil
 }
 
 func (uvm *UtilityVM) StopLogForwarding(ctx context.Context) error {
-	// Implementation for stopping the log forwarder
-	if uvm.OS() == "windows" && uvm.gc != nil {
-		wcaps := gcs.GetWCOWCapabilities(uvm.gc.Capabilities())
-		if wcaps != nil && wcaps.IsLogForwardingSupported() {
-			req := guestrequest.LogForwardServiceRPCRequest{
-				RPCType:  guestrequest.RPCStopLogForwarding,
-				Settings: "",
-			}
-			err := uvm.gc.ModifyServiceSettings(ctx, prot.LogForwardService, req)
-			if err != nil {
-				return err
-			}
+	// Implementation for stopping the log forwarding service
+	if uvm.OS() != "windows" || uvm.gc == nil {
+		return errNotSupported
+	}
+
+	wcaps := gcs.GetWCOWCapabilities(uvm.gc.Capabilities())
+	if wcaps != nil && wcaps.IsLogForwardingSupported() {
+		req := guestrequest.LogForwardServiceRPCRequest{
+			RPCType:  guestrequest.RPCStopLogForwarding,
+			Settings: "",
+		}
+		err := uvm.gc.ModifyServiceSettings(ctx, prot.LogForwardService, req)
+		if err != nil {
+			return err
 		}
 	}
 	return nil
 }
 
 func (uvm *UtilityVM) SetLogSources(ctx context.Context) error {
 	// Implementation for setting the log sources
-	if uvm.OS() == "windows" && uvm.logSources != "" && uvm.gc != nil {
-		wcaps := gcs.GetWCOWCapabilities(uvm.gc.Capabilities())
-		if wcaps != nil && wcaps.IsLogForwardingSupported() {
-			// Make a call to the GCS to set the ETW providers
-			req := guestrequest.LogForwardServiceRPCRequest{
-				RPCType:  guestrequest.RPCModifyServiceSettings,
-				Settings: uvm.logSources,
-			}
-			err := uvm.gc.ModifyServiceSettings(ctx, prot.LogForwardService, req)
-			if err != nil {
-				return err
-			}
+	if uvm.OS() != "windows" || uvm.gc == nil {
+		return errNotSupported
+	}
+
+	wcaps := gcs.GetWCOWCapabilities(uvm.gc.Capabilities())
+	if wcaps != nil && wcaps.IsLogForwardingSupported() {
+		// Make a call to the GCS to set the ETW providers
+		req := guestrequest.LogForwardServiceRPCRequest{
+			RPCType:  guestrequest.RPCModifyServiceSettings,
+			Settings: uvm.logSources,
+		}
+		err := uvm.gc.ModifyServiceSettings(ctx, prot.LogForwardService, req)
+		if err != nil {
+			return err
 		}
 	}
 	return nil

internal/uvm/start.go

@@ -70,13 +70,11 @@ func (e *gcsLogEntry) UnmarshalJSON(b []byte) error {
 	if e.Fields["Source"] == "ETW" {
 		// Windows ETW log entry
 		// Original ETW Event Data may have "message" or "Message" field instead of "msg"
-
-		if e.Message == "" && e.Fields["message"] != nil {
-			e.Message = e.Fields["message"].(string)
+		if msg, ok := e.Fields["message"].(string); ok {
+			e.Message = msg
 			delete(e.Fields, "message")
-		}
-		if e.Message == "" && e.Fields["Message"] != nil {
-			e.Message = e.Fields["Message"].(string)
+		} else if msg, ok := e.Fields["Message"].(string); ok {
+			e.Message = msg
 			delete(e.Fields, "Message")
 		}
 	}
@@ -220,33 +218,31 @@ func (uvm *UtilityVM) Start(ctx context.Context) (err error) {
 			// For windows, the Listener can recieve a connection later, so we
 			// start the output handler in a goroutine with a non-timeout context.
 			// This allows the output handler to run independently of the UVM Create's
-			// lifecycle. This approach potentially allows to wait for reconnections,
+			// lifecycle. The approach potentially allows to wait for reconnections too,
 			// while limiting the number of concurrent connections to 1.
 			// This is useful for the case when logging service is restarted.
-			g.Go(func() error {
-				go func() {
-					var wg sync.WaitGroup
-					uvm.outputListener = netutil.LimitListener(uvm.outputListener, 1)
-					for {
-						conn, err := uvm.accept(context.Background(), uvm.outputListener, false)
-						if err != nil {
-							e.WithError(err).Error("failed to connect to log socket")
-							close(uvm.outputProcessingDone)
-							break
-						}
-						wg.Add(1)
-						go func() {
-							defer wg.Done()
-							e.Info("uvm output handler starting")
-							uvm.outputHandler(conn)
-						}()
-						e.Info("uvm output handler finished")
+			go func() {
+				var wg sync.WaitGroup
+				uvm.outputListener = netutil.LimitListener(uvm.outputListener, 1)
+				for {
+					conn, err := uvm.accept(context.WithoutCancel(ctx), uvm.outputListener, false)
+					if err != nil {
+						e.WithError(err).Error("failed to connect to log socket")
+						break
 					}
-					wg.Wait()
+					wg.Add(1)
+					go func() {
+						defer wg.Done()
+						e.Info("uvm output handler starting")
+						uvm.outputHandler(conn)
+					}()
+					e.Info("uvm output handler finished")
+				}
+				wg.Wait()
+				if _, ok := <-uvm.outputProcessingDone; ok {
 					close(uvm.outputProcessingDone)
-				}()
-				return nil
-			})
+				}
+			}()
 		default:
 			// Default handling
 			g.Go(func() error {

pkg/annotations/annotations.go

@@ -401,7 +401,31 @@ const (
 	// VSMBNoDirectMap specifies that no direct mapping should be used for any VSMBs added to the UVM.
 	VSMBNoDirectMap = "io.microsoft.virtualmachine.wcow.virtualSMB.nodirectmap"
 
-	// LogSources specifies the ETW providers to be set for the logging service.
+	// LogSources specifies the ETW providers to be set for the logging service as a base64-encoded JSON string.
+	//
+	// For example:
+	//
+	//	{
+	//	  "logConfig": {
+	//	    "sources": [
+	//	      {
+	//	        "type": "ETW",
+	//	        "providers": [
+	//	          {
+	//	            "providerGuid": "80CE50DE-D264-4581-950D-ABADEEE0D340",
+	//	            "providerName": "Microsoft.Windows.HyperV.Compute",
+	//	            "level": "Information"
+	//	          }
+	//	        ]
+	//	      }
+	//	    ]
+	//	  }
+	//	}
+	//
+	// Would be encoded as:
+	//
+	// 	"io.microsoft.virtualmachine.wcow.logsources" =
+	// 		"eyJsb2dDb25maWciOnsic291cmNlcyI6W3sidHlwZSI6IkVUVyIsInByb3ZpZGVycyI6W3sicHJvdmlkZXJHdWlkIjoiODBDRTUwREUtRDI2NC00NTgxLTk1MEQtQUJBREVFRTBEMzQwIiwicHJvdmlkZXJOYW1lIjoiTWljcm9zb2Z0LldpbmRvd3MuSHlwZXJWLkNvbXB1dGUiLCJsZXZlbCI6IkluZm9ybWF0aW9uIn1dfV19fQ=="
 	LogSources = "io.microsoft.virtualmachine.wcow.logsources"
 
 	// ForwardLogs specifies whether to forward logs to the host or not.