CSP Parsing with Golang

go.mod

module gist.github.com/rzhade3/88a5ab6d8c2fc6016cad0a2b507a88a4

go 1.20

main.go

package main

import (
	"reflect"
	"strings"
)

// Represents a Content-Security-Policy header
type CSP struct {
	DefaultSrc     string
	ScriptSrc      string
	StyleSrc       string
	ImgSrc         string
	ConnectSrc     string
	FontSrc        string
	ObjectSrc      string
	MediaSrc       string
	FrameSrc       string
	FrameAncestors string
	Sandbox        string
	// For other headers, just put them in a map
	Other map[string]string
}

// Maps the CSP directive to the corresponding field in the CSP struct
var directiveToField = map[string]string{
	"default-src":     "DefaultSrc",
	"script-src":      "ScriptSrc",
	"style-src":       "StyleSrc",
	"img-src":         "ImgSrc",
	"connect-src":     "ConnectSrc",
	"font-src":        "FontSrc",
	"object-src":      "ObjectSrc",
	"media-src":       "MediaSrc",
	"frame-src":       "FrameSrc",
	"frame-ancestors": "FrameAncestors",
	"sandbox":         "Sandbox",
}

func ParseHeader(header string) CSP {
	directives := strings.Split(header, ";")
	csp := CSP{Other: make(map[string]string)}
	for _, directive := range directives {
		directive = strings.TrimSpace(directive)
		directiveName := strings.SplitN(strings.TrimSpace(directive), " ", 2)[0]
		directiveValue := strings.TrimSpace(strings.SplitN(directive, " ", 2)[1])
		if val, ok := directiveToField[directiveName]; ok {
			// If the directive is a known one, put the value in the corresponding field
			// Reflect is used to set the field value
			// See https://stackoverflow.com/questions/6395076/using-a-variable-as-a-struct-field-name-in-go
			reflect.ValueOf(&csp).Elem().FieldByName(val).SetString(directiveValue)
		} else {
			// For other headers, just put them in a map
			csp.Other[directiveName] = directiveValue
		}
	}
	return csp
}

func SetHeader(csp CSP) string {
	directives := []string{}
	directives = append(directives, "default-src "+csp.DefaultSrc)
	directives = append(directives, "script-src "+csp.ScriptSrc)
	directives = append(directives, "style-src "+csp.StyleSrc)
	directives = append(directives, "img-src "+csp.ImgSrc)
	directives = append(directives, "connect-src "+csp.ConnectSrc)
	directives = append(directives, "font-src "+csp.FontSrc)
	directives = append(directives, "object-src "+csp.ObjectSrc)
	directives = append(directives, "media-src "+csp.MediaSrc)
	directives = append(directives, "frame-src "+csp.FrameSrc)
	directives = append(directives, "frame-ancestors "+csp.FrameAncestors)
	directives = append(directives, "sandbox "+csp.Sandbox)
	for directive, value := range csp.Other {
		directives = append(directives, directive+" "+value)
	}
	return strings.Join(directives, "; ")
}

main_test.go

package main

import (
	"testing"
)

func TestParseHeader(t *testing.T) {
	test_csp_header := "default-src 'self'; script-src 'self' 'unsafe-inline' 'unsafe-eval'; style-src 'self' 'unsafe-inline'; img-src 'self' data:; connect-src 'self'; font-src 'self'; object-src 'none'; media-src 'none'; frame-src 'none'; frame-ancestors 'none'; sandbox allow-forms allow-same-origin allow-scripts; report-uri /foo"
	csp := ParseHeader(test_csp_header)
	// Assert that the CSP is parsed correctly
	if csp.DefaultSrc != "'self'" {
		t.Errorf("Expected default-src to be 'self', got %s", csp.DefaultSrc)
	}
	if csp.ScriptSrc != "'self' 'unsafe-inline' 'unsafe-eval'" {
		t.Errorf("Expected script-src to be 'self' 'unsafe-inline' 'unsafe-eval', got %s", csp.ScriptSrc)
	}
	if csp.StyleSrc != "'self' 'unsafe-inline'" {
		t.Errorf("Expected style-src to be 'self' 'unsafe-inline', got %s", csp.StyleSrc)
	}
	if csp.ImgSrc != "'self' data:" {
		t.Errorf("Expected img-src to be 'self' data:, got %s", csp.ImgSrc)
	}
	if csp.ConnectSrc != "'self'" {
		t.Errorf("Expected connect-src to be 'self', got %s", csp.ConnectSrc)
	}
	if csp.FontSrc != "'self'" {
		t.Errorf("Expected font-src to be 'self', got %s", csp.FontSrc)
	}
	if csp.ObjectSrc != "'none'" {
		t.Errorf("Expected object-src to be 'none', got %s", csp.ObjectSrc)
	}
	if csp.MediaSrc != "'none'" {
		t.Errorf("Expected media-src to be 'none', got %s", csp.MediaSrc)
	}
	if csp.FrameSrc != "'none'" {
		t.Errorf("Expected frame-src to be 'none', got %s", csp.FrameSrc)
	}
	if csp.FrameAncestors != "'none'" {
		t.Errorf("Expected frame-ancestors to be 'none', got %s", csp.FrameAncestors)
	}
	if csp.Sandbox != "allow-forms allow-same-origin allow-scripts" {
		t.Errorf("Expected sandbox to be allow-forms allow-same-origin allow-scripts, got %s", csp.Sandbox)
	}
	if csp.Other["report-uri"] != "/foo" {
		t.Errorf("Expected report-uri to be /foo, got %s", csp.Other["report-uri"])
	}
}

func TestSetHeader(t *testing.T) {
	dummy_csp := CSP{
		DefaultSrc:     "'self'",
		ScriptSrc:      "'self' 'unsafe-inline' 'unsafe-eval'",
		StyleSrc:       "'self' 'unsafe-inline'",
		ImgSrc:         "'self' data:",
		ConnectSrc:     "'self'",
		FontSrc:        "'self'",
		ObjectSrc:      "'none'",
		MediaSrc:       "'none'",
		FrameSrc:       "'none'",
		FrameAncestors: "'none'",
		Sandbox:        "allow-forms allow-same-origin allow-scripts",
		Other:          map[string]string{"report-uri": "/foo"},
	}
	csp_header := SetHeader(dummy_csp)
	// Assert that the CSP is set correctly
	if csp_header != "default-src 'self'; script-src 'self' 'unsafe-inline' 'unsafe-eval'; style-src 'self' 'unsafe-inline'; img-src 'self' data:; connect-src 'self'; font-src 'self'; object-src 'none'; media-src 'none'; frame-src 'none'; frame-ancestors 'none'; sandbox allow-forms allow-same-origin allow-scripts; report-uri /foo" {
		t.Errorf("Expected CSP header to be default-src 'self'; script-src 'self' 'unsafe-inline' 'unsafe-eval'; style-src 'self' 'unsafe-inline'; img-src 'self' data:; connect-src 'self'; font-src 'self'; object-src 'none'; media-src 'none'; frame-src 'none'; frame-ancestors 'none'; sandbox allow-forms allow-same-origin allow-scripts; report-uri /foo, got %s", csp_header)
	}
}