Varied bugfixes

client/incus.go

@@ -154,6 +154,17 @@ func (r *ProtocolIncus) DoHTTP(req *http.Request) (*http.Response, error) {
 	return r.http.Do(req)
 }
 
+// DoWebsocket performs a websocket connection, using OIDC authentication if set.
+func (r *ProtocolIncus) DoWebsocket(dialer websocket.Dialer, uri string, req *http.Request) (*websocket.Conn, *http.Response, error) {
+	r.addClientHeaders(req)
+
+	if r.oidcClient != nil {
+		return r.oidcClient.dial(dialer, uri, req)
+	}
+
+	return dialer.Dial(uri, req.Header)
+}
+
 // addClientHeaders sets headers from client settings.
 // User-Agent (if r.httpUserAgent is set).
 // X-Incus-authenticated (if r.requireAuthenticated is set).
@@ -436,10 +447,9 @@ func (r *ProtocolIncus) rawWebsocket(url string) (*websocket.Conn, error) {
 	// Create temporary http.Request using the http url, not the ws one, so that we can add the client headers
 	// for the websocket request.
 	req := &http.Request{URL: &r.httpBaseURL, Header: http.Header{}}
-	r.addClientHeaders(req)
 
 	// Establish the connection
-	conn, resp, err := dialer.Dial(url, req.Header)
+	conn, resp, err := r.DoWebsocket(dialer, url, req)
 	if err != nil {
 		if resp != nil {
 			_, _, err = incusParseResponse(resp)

client/incus_oidc.go

@@ -13,6 +13,7 @@ import (
 	"syscall"
 	"time"
 
+	"github.com/gorilla/websocket"
 	"github.com/zitadel/oidc/v3/pkg/client/rp"
 	httphelper "github.com/zitadel/oidc/v3/pkg/http"
 	"github.com/zitadel/oidc/v3/pkg/oidc"
@@ -147,6 +148,40 @@ func (o *oidcClient) do(req *http.Request) (*http.Response, error) {
 	return resp, nil
 }
 
+// dial function executes a websocket request and handles OIDC authentication and refresh.
+func (o *oidcClient) dial(dialer websocket.Dialer, uri string, req *http.Request) (*websocket.Conn, *http.Response, error) {
+	conn, resp, err := dialer.Dial(uri, req.Header)
+	if err != nil && resp == nil {
+		return nil, nil, err
+	}
+
+	// Return immediately if the error is not HTTP status unauthorized.
+	if conn != nil && resp.StatusCode != http.StatusUnauthorized {
+		return conn, resp, nil
+	}
+
+	issuer := resp.Header.Get("X-Incus-OIDC-issuer")
+	clientID := resp.Header.Get("X-Incus-OIDC-clientid")
+	audience := resp.Header.Get("X-Incus-OIDC-audience")
+
+	if issuer == "" || clientID == "" {
+		return nil, resp, err
+	}
+
+	err = o.refresh(issuer, clientID)
+	if err != nil {
+		err = o.authenticate(issuer, clientID, audience)
+		if err != nil {
+			return nil, resp, err
+		}
+	}
+
+	// Set the new access token in the header.
+	req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", o.tokens.AccessToken))
+
+	return dialer.Dial(uri, req.Header)
+}
+
 // getProvider initializes a new OpenID Connect Relying Party for a given issuer and clientID.
 // The function also creates a secure CookieHandler with random encryption and hash keys, and applies a series of configurations on the Relying Party.
 func (o *oidcClient) getProvider(issuer string, clientID string) (rp.RelyingParty, error) {