Merge pull request #66 from zitadel/redirect-loop

fix: handle redirect loop

pkg/provider/logout.go

@@ -44,7 +44,7 @@ func (p *IdentityProvider) logoutHandleFunc(w http.ResponseWriter, r *http.Reque
 			return nil
 		},
 		func() {
-			http.Error(w, fmt.Errorf("failed to parse form: %w", err).Error(), http.StatusInternalServerError)
+			response.sendBackLogoutResponse(w, response.makeDeniedLogoutResponse(fmt.Errorf("failed to parse form: %w", err).Error(), p.timeFormat))
 		},
 	)
 
@@ -60,7 +60,7 @@ func (p *IdentityProvider) logoutHandleFunc(w http.ResponseWriter, r *http.Reque
 			return nil
 		},
 		func() {
-			response.sendBackLogoutResponse(w, response.makeUnsupportedlLogoutResponse(fmt.Errorf("failed to decode request: %w", err).Error(), p.timeFormat))
+			response.sendBackLogoutResponse(w, response.makeDeniedLogoutResponse(fmt.Errorf("failed to decode request: %w", err).Error(), p.timeFormat))
 		},
 	)
 

pkg/provider/logout_response.go

@@ -1,14 +1,12 @@
 package provider
 
 import (
-	"bufio"
-	"bytes"
 	"encoding/base64"
-	"encoding/xml"
 	"html/template"
 	"net/http"
 	"time"
 
+	"github.com/zitadel/saml/pkg/provider/xml"
 	"github.com/zitadel/saml/pkg/provider/xml/saml"
 	"github.com/zitadel/saml/pkg/provider/xml/samlp"
 )
@@ -31,33 +29,23 @@ type LogoutResponseForm struct {
 }
 
 func (r *LogoutResponse) sendBackLogoutResponse(w http.ResponseWriter, resp *samlp.LogoutResponseType) {
-	var xmlbuff bytes.Buffer
-
-	memWriter := bufio.NewWriter(&xmlbuff)
-	_, err := memWriter.Write([]byte(xml.Header))
+	respData, err := xml.Marshal(resp)
 	if err != nil {
 		r.ErrorFunc(err)
 		return
 	}
 
-	encoder := xml.NewEncoder(memWriter)
-	err = encoder.Encode(resp)
-	if err != nil {
-		r.ErrorFunc(err)
+	if r.LogoutURL == "" {
+		if err := xml.Write(w, respData); err != nil {
+			r.ErrorFunc(err)
+			return
+		}
 		return
 	}
 
-	err = memWriter.Flush()
-	if err != nil {
-		r.ErrorFunc(err)
-		return
-	}
-
-	samlMessage := base64.StdEncoding.EncodeToString(xmlbuff.Bytes())
-
 	data := LogoutResponseForm{
 		RelayState:   r.RelayState,
-		SAMLResponse: samlMessage,
+		SAMLResponse: base64.StdEncoding.EncodeToString(respData),
 		LogoutURL:    r.LogoutURL,
 	}
 

pkg/provider/redirect.go

@@ -71,12 +71,12 @@ func createRedirectSignature(
 	idp *IdentityProvider,
 	response *Response,
 ) error {
-	respStr, err := xml.Marshal(samlResponse)
+	resp, err := xml.Marshal(samlResponse)
 	if err != nil {
 		return err
 	}
 
-	respData, err := xml.DeflateAndBase64([]byte(respStr))
+	respData, err := xml.DeflateAndBase64(resp)
 	if err != nil {
 		return err
 	}

pkg/provider/response.go

@@ -41,16 +41,37 @@ type Response struct {
 }
 
 func (r *Response) doResponse(request *http.Request, w http.ResponseWriter, response string) {
+
+}
+
+type AuthResponseForm struct {
+	RelayState                  string
+	SAMLResponse                string
+	AssertionConsumerServiceURL string
+}
+
+func (r *Response) sendBackResponse(
+	req *http.Request,
+	w http.ResponseWriter,
+	resp *samlp.ResponseType,
+) {
+	respData, err := xml.Marshal(resp)
+	if err != nil {
+		r.ErrorFunc(err)
+		return
+	}
+
 	if r.AcsUrl == "" {
-		if err := xml.Write(w, []byte(response)); err != nil {
+		if err := xml.Write(w, respData); err != nil {
 			r.ErrorFunc(err)
 			return
 		}
+		return
 	}
 
 	switch r.ProtocolBinding {
 	case PostBinding:
-		respData := base64.StdEncoding.EncodeToString([]byte(response))
+		respData := base64.StdEncoding.EncodeToString(respData)
 
 		data := AuthResponseForm{
 			r.RelayState,
@@ -63,39 +84,19 @@ func (r *Response) doResponse(request *http.Request, w http.ResponseWriter, resp
 			return
 		}
 	case RedirectBinding:
-		respData, err := xml.DeflateAndBase64([]byte(response))
+		respData, err := xml.DeflateAndBase64(respData)
 		if err != nil {
 			r.ErrorFunc(err)
 			return
 		}
 
-		http.Redirect(w, request, fmt.Sprintf("%s?%s", r.AcsUrl, buildRedirectQuery(string(respData), r.RelayState, r.SigAlg, r.Signature)), http.StatusFound)
+		http.Redirect(w, req, fmt.Sprintf("%s?%s", r.AcsUrl, buildRedirectQuery(string(respData), r.RelayState, r.SigAlg, r.Signature)), http.StatusFound)
 		return
 	default:
 		//TODO: no binding
 	}
 }
 
-type AuthResponseForm struct {
-	RelayState                  string
-	SAMLResponse                string
-	AssertionConsumerServiceURL string
-}
-
-func (r *Response) sendBackResponse(
-	req *http.Request,
-	w http.ResponseWriter,
-	resp *samlp.ResponseType,
-) {
-	respStr, err := xml.Marshal(resp)
-	if err != nil {
-		r.ErrorFunc(err)
-		return
-	}
-
-	r.doResponse(req, w, respStr)
-}
-
 func (r *Response) makeUnsupportedBindingResponse(
 	message string,
 	timeFormat string,

pkg/provider/signature/signature_test.go

@@ -512,7 +512,7 @@ func TestSignature_CreatePost(t *testing.T) {
 			}
 			resp.Signature = sig
 
-			respStr, err := saml_xml.Marshal(resp)
+			respData, err := saml_xml.Marshal(resp)
 			if err != nil {
 				if (err != nil) != tt.res.err {
 					t.Errorf("Create() marshall response for signing")
@@ -521,7 +521,7 @@ func TestSignature_CreatePost(t *testing.T) {
 			}
 
 			doc := etree.NewDocument()
-			if err := doc.ReadFromBytes([]byte(respStr)); err != nil {
+			if err := doc.ReadFromBytes(respData); err != nil {
 				if (err != nil) != tt.res.err {
 					t.Errorf("Cert() failed to read response")
 				}

pkg/provider/sso.go

@@ -61,7 +61,7 @@ func (p *IdentityProvider) ssoHandleFunc(w http.ResponseWriter, r *http.Request)
 			return nil
 		},
 		func() {
-			http.Error(w, fmt.Errorf("failed to parse form: %w", err).Error(), http.StatusInternalServerError)
+			response.sendBackResponse(r, w, response.makeDeniedResponse(fmt.Errorf("failed to parse form").Error(), p.timeFormat))
 		},
 	)
 

pkg/provider/sso_test.go

@@ -559,8 +559,8 @@ func TestSSO_ssoHandleFunc(t *testing.T) {
 				},
 			},
 			res{
-				code:  500,
-				state: "",
+				code:  200,
+				state: StatusCodeRequestDenied,
 				err:   false,
 			}},
 		{

pkg/provider/xml/xml.go

@@ -19,27 +19,27 @@ const (
 	EncodingDeflate = "urn:oasis:names:tc:SAML:2.0:bindings:URL-Encoding:DEFLATE"
 )
 
-func Marshal(data interface{}) (string, error) {
+func Marshal(data interface{}) ([]byte, error) {
 	var xmlbuff bytes.Buffer
 
 	memWriter := bufio.NewWriter(&xmlbuff)
 	_, err := memWriter.Write([]byte(xml.Header))
 	if err != nil {
-		return "", err
+		return nil, err
 	}
 
 	encoder := xml.NewEncoder(memWriter)
 	err = encoder.Encode(data)
 	if err != nil {
-		return "", err
+		return nil, err
 	}
 
 	err = memWriter.Flush()
 	if err != nil {
-		return "", err
+		return nil, err
 	}
 
-	return xmlbuff.String(), nil
+	return xmlbuff.Bytes(), nil
 }
 
 func DeflateAndBase64(data []byte) ([]byte, error) {

pkg/provider/xml/xml_test.go

@@ -1,6 +1,7 @@
 package xml_test
 
 import (
+	"reflect"
 	"testing"
 
 	"github.com/zitadel/saml/pkg/provider/xml"
@@ -12,7 +13,7 @@ type XML struct {
 
 func Test_XmlMarshal(t *testing.T) {
 	type res struct {
-		metadata string
+		metadata []byte
 		err      bool
 	}
 
@@ -25,23 +26,23 @@ func Test_XmlMarshal(t *testing.T) {
 			name: "xml struct",
 			arg:  "<test></test>",
 			res: res{
-				metadata: "<?xml version=\"1.0\" encoding=\"UTF-8\"?>\n<XML><test></test></XML>",
+				metadata: []byte("<?xml version=\"1.0\" encoding=\"UTF-8\"?>\n<XML><test></test></XML>"),
 				err:      false,
 			},
 		},
 	}
-
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			xmlStruct := XML{InnerXml: tt.arg}
 
-			xmlStr, err := xml.Marshal(xmlStruct)
+			xmlData, err := xml.Marshal(xmlStruct)
 			if (err != nil) != tt.res.err {
 				t.Errorf("Marshal() error: %v", err)
 				return
 			}
-			if xmlStr != tt.res.metadata {
-				t.Errorf("Marshal() error expected: %v, got %v", tt.res.metadata, xmlStr)
+
+			if !reflect.DeepEqual(tt.res.metadata, xmlData) {
+				t.Errorf("Marshal() error expected: %v, got %v", tt.res.metadata, xmlData)
 				return
 			}
 		})