From b243c7015755e697c2a2238b95fced9c083d96ed Mon Sep 17 00:00:00 2001
From: Leandro Dorileo <dorileo@google.com>
Date: Mon, 25 Sep 2023 10:49:00 -0700
Subject: [PATCH] oslogin: introduce a feature flag to cert auth (#298)

With this flag customers can choose to disable the configuration
management for OS Login cert based authentication.
---
 google_guest_agent/cfg/cfg.go      | 11 +++++++++++
 google_guest_agent/oslogin.go      | 13 +++++++++----
 google_guest_agent/oslogin_test.go |  5 +++++
 instance_configs.cfg               |  3 +++
 4 files changed, 28 insertions(+), 4 deletions(-)

diff --git a/google_guest_agent/cfg/cfg.go b/google_guest_agent/cfg/cfg.go
index bcc56c103..b53071eda 100644
--- a/google_guest_agent/cfg/cfg.go
+++ b/google_guest_agent/cfg/cfg.go
@@ -84,6 +84,9 @@ dhcp_command =
 ip_forwarding = true
 setup = true
 
+[OSLogin]
+cert_authentication = true
+
 [Snapshots]
 enabled = false
 snapshot_service_ip = 169.254.169.254
@@ -139,6 +142,9 @@ type Sections struct {
 	// as well as the commands definitions for network configuration.
 	NetworkInterfaces *NetworkInterfaces `ini:"NetworkInterfaces,omitempty"`
 
+	// OSLogin defines the OS Login configuration options.
+	OSLogin *OSLogin `ini:"OSLogin,omitempty"`
+
 	// Snpashots defines the snapshot listener configuration and behavior i.e. the server address and port.
 	Snapshots *Snapshots `ini:"Snapshots,omitempty"`
 
@@ -227,6 +233,11 @@ type MetadataScripts struct {
 	SysprepSpecialize bool   `ini:"sysprep_specialize,omitempty"`
 }
 
+// OSLogin contains the configurations of OSLogin section.
+type OSLogin struct {
+	CertAuthentication bool `ini:"cert_authentication,omitempty"`
+}
+
 // NetworkInterfaces contains the configurations of NetworkInterfaces section.
 type NetworkInterfaces struct {
 	DHCPCommand  string `ini:"dhcp_command,omitempty"`
diff --git a/google_guest_agent/oslogin.go b/google_guest_agent/oslogin.go
index 234405a4b..af8c9d548 100644
--- a/google_guest_agent/oslogin.go
+++ b/google_guest_agent/oslogin.go
@@ -208,10 +208,10 @@ func updateSSHConfig(sshConfig string, enable, twofactor, pamlessAuthStack, skey
 		}
 	}
 	authorizedKeysUser := "AuthorizedKeysCommandUser root"
+
+	// Certificate based authentication.
 	authorizedPrincipalsCommand := "AuthorizedPrincipalsCommand /usr/bin/google_authorized_principals %u %k"
 	authorizedPrincipalsUser := "AuthorizedPrincipalsCommandUser root"
-
-	// TODO: only enable this key configuration if certs mechanism is enabled
 	trustedUserCAKeys := "TrustedUserCAKeys " + sshtrustedca.DefaultPipePath
 
 	twoFactorAuthMethods := "AuthenticationMethods publickey,keyboard-interactive"
@@ -225,8 +225,13 @@ func updateSSHConfig(sshConfig string, enable, twofactor, pamlessAuthStack, skey
 	filtered := filterGoogleLines(string(sshConfig))
 
 	if enable {
-		osLoginBlock := []string{googleBlockStart, authorizedKeysCommand, authorizedKeysUser, trustedUserCAKeys}
-		if pamlessAuthStack {
+		osLoginBlock := []string{googleBlockStart, authorizedKeysCommand, authorizedKeysUser}
+
+		if cfg.Get().OSLogin.CertAuthentication {
+			osLoginBlock = append(osLoginBlock, trustedUserCAKeys)
+		}
+
+		if pamlessAuthStack && cfg.Get().OSLogin.CertAuthentication {
 			osLoginBlock = append(osLoginBlock, authorizedPrincipalsCommand, authorizedPrincipalsUser)
 		}
 
diff --git a/google_guest_agent/oslogin_test.go b/google_guest_agent/oslogin_test.go
index 03b28f427..27aa9078b 100644
--- a/google_guest_agent/oslogin_test.go
+++ b/google_guest_agent/oslogin_test.go
@@ -20,6 +20,7 @@ import (
 	"strings"
 	"testing"
 
+	"github.com/GoogleCloudPlatform/guest-agent/google_guest_agent/cfg"
 	"github.com/GoogleCloudPlatform/guest-agent/google_guest_agent/events/sshtrustedca"
 	"github.com/GoogleCloudPlatform/guest-agent/metadata"
 )
@@ -310,6 +311,10 @@ func TestUpdateSSHConfig(t *testing.T) {
 		},
 	}
 
+	if err := cfg.Load(nil); err != nil {
+		t.Fatalf("Failed to initialize configuration manager: %+v", err)
+	}
+
 	for idx, tt := range tests {
 		contents := strings.Join(tt.contents, "\n")
 		want := strings.Join(tt.want, "\n")
diff --git a/instance_configs.cfg b/instance_configs.cfg
index b515cf8e2..301c2deb3 100644
--- a/instance_configs.cfg
+++ b/instance_configs.cfg
@@ -36,3 +36,6 @@ dhclient_script = /sbin/google-dhclient-script
 dhcp_command = 
 ip_forwarding = true
 setup = true
+
+[OSLogin]
+cert_authentication = true