package authres

import (
        "sort"
        "strings"
        "unicode"
)

// Format formats an Authentication-Results header.
func Format(identity string, results []Result) string {
        s := identity

        if len(results) == 0 {
                s += "; none"
                return s
        }

        for _, r := range results {
                method := resultMethod(r)
                value, params := r.format()

                s += "; " + method + "=" + string(value) + " " + formatParams(params)
        }

        return s
}

func resultMethod(r Result) string {
        switch r := r.(type) {
        case *AuthResult:
                return "auth"
        case *DKIMResult:
                return "dkim"
        case *DomainKeysResult:
                return "domainkeys"
        case *IPRevResult:
                return "iprev"
        case *SenderIDResult:
                return "sender-id"
        case *SPFResult:
                return "spf"
        case *DMARCResult:
                return "dmarc"
        case *GenericResult:
                return r.Method
        default:
                return ""
        }
}

func formatParams(params map[string]string) string {
        keys := make([]string, 0, len(params))
        for k := range params {
                if k == "reason" {
                        continue
                }
                keys = append(keys, k)
        }
        sort.Strings(keys)
        if params["reason"] != "" {
                keys = append([]string{"reason"}, keys...)
        }

        s := ""
        i := 0
        for _, k := range keys {
                if params[k] == "" {
                        continue
                }

                if i > 0 {
                        s += " "
                }

                var value string
                if k == "reason" {
                        value = formatValue(params[k])
                } else {
                        value = formatPvalue(params[k])
                }
                s += k + "=" + value
                i++
        }

        return s
}

var tspecials = map[rune]struct{}{
        '(': {}, ')': {}, '<': {}, '>': {}, '@': {},
        ',': {}, ';': {}, ':': {}, '\\': {}, '"': {},
        '/': {}, '[': {}, ']': {}, '?': {}, '=': {},
}

func formatValue(s string) string {
        // value := token / quoted-string
        // token := 1*<any (US-ASCII) CHAR except SPACE, CTLs,
        //            or tspecials>
        // tspecials :=  "(" / ")" / "<" / ">" / "@" /
        //               "," / ";" / ":" / "\" / <">
        //               "/" / "[" / "]" / "?" / "="
        //               ; Must be in quoted-string,
        //               ; to use within parameter values

        shouldQuote := false
        for _, ch := range s {
                if _, special := tspecials[ch]; ch <= ' ' /* SPACE or CTL */ || special {
                        shouldQuote = true
                }
        }

        if shouldQuote {
                return `"` + strings.Replace(s, `"`, `\"`, -1) + `"`
        }
        return s
}

var addressOk = map[rune]struct{}{
        // Most ASCII punctuation except for:
        //  ( ) = "
        // as these can cause issues due to ambiguous ABNF rules.
        // I.e. technically mentioned characters can be left unquoted, but they can
        // be interpreted as parts of non-quoted parameters or comments so it is
        // better to quote them.
        '#': {}, '$': {}, '%': {}, '&': {},
        '\'': {}, '*': {}, '+': {}, ',': {},
        '.': {}, '/': {}, '-': {}, '@': {},
        '[': {}, ']': {}, '\\': {}, '^': {},
        '_': {}, '`': {}, '{': {}, '|': {},
        '}': {}, '~': {},
}

func formatPvalue(s string) string {
        // pvalue = [CFWS] ( value / [ [ local-part ] "@" ] domain-name )
        //          [CFWS]

        // Experience shows that implementers often "forget" that things can
        // be quoted in various places where they are usually not quoted
        // so we can't get away by just quoting everything.

        // Relevant ABNF rules are much complicated than that, but this
        // will catch most of the cases and we can fallback to quoting
        // for others.
        addressLike := true
        for _, ch := range s {
                if _, ok := addressOk[ch]; !unicode.IsLetter(ch) && !unicode.IsDigit(ch) && !ok {
                        addressLike = false
                }
        }

        if addressLike {
                return s
        }
        return formatValue(s)
}

package authres

import (
        "errors"
        "fmt"
        "strconv"
        "strings"
        "unicode"
)

// ResultValue is an authentication result value, as defined in RFC 5451 section
// 6.3.
type ResultValue string

const (
        ResultNone      ResultValue = "none"
        ResultPass                  = "pass"
        ResultFail                  = "fail"
        ResultPolicy                = "policy"
        ResultNeutral               = "neutral"
        ResultTempError             = "temperror"
        ResultPermError             = "permerror"
        ResultHardFail              = "hardfail"
        ResultSoftFail              = "softfail"
)

// Result is an authentication result.
type Result interface {
        parse(value ResultValue, params map[string]string) error
        format() (value ResultValue, params map[string]string)
}

type AuthResult struct {
        Value  ResultValue
        Reason string
        Auth   string
}

func (r *AuthResult) parse(value ResultValue, params map[string]string) error {
        r.Value = value
        r.Reason = params["reason"]
        r.Auth = params["smtp.auth"]
        return nil
}

func (r *AuthResult) format() (ResultValue, map[string]string) {
        return r.Value, map[string]string{"smtp.auth": r.Auth}
}

type DKIMResult struct {
        Value      ResultValue
        Reason     string
        Domain     string
        Identifier string
}

func (r *DKIMResult) parse(value ResultValue, params map[string]string) error {
        r.Value = value
        r.Reason = params["reason"]
        r.Domain = params["header.d"]
        r.Identifier = params["header.i"]
        return nil
}

func (r *DKIMResult) format() (ResultValue, map[string]string) {
        return r.Value, map[string]string{
                "reason":   r.Reason,
                "header.d": r.Domain,
                "header.i": r.Identifier,
        }
}

type DomainKeysResult struct {
        Value  ResultValue
        Reason string
        Domain string
        From   string
        Sender string
}

func (r *DomainKeysResult) parse(value ResultValue, params map[string]string) error {
        r.Value = value
        r.Reason = params["reason"]
        r.Domain = params["header.d"]
        r.From = params["header.from"]
        r.Sender = params["header.sender"]
        return nil
}

func (r *DomainKeysResult) format() (ResultValue, map[string]string) {
        return r.Value, map[string]string{
                "reason":        r.Reason,
                "header.d":      r.Domain,
                "header.from":   r.From,
                "header.sender": r.Sender,
        }
}

type IPRevResult struct {
        Value  ResultValue
        Reason string
        IP     string
}

func (r *IPRevResult) parse(value ResultValue, params map[string]string) error {
        r.Value = value
        r.Reason = params["reason"]
        r.IP = params["policy.iprev"]
        return nil
}

func (r *IPRevResult) format() (ResultValue, map[string]string) {
        return r.Value, map[string]string{
                "reason":       r.Reason,
                "policy.iprev": r.IP,
        }
}

type SenderIDResult struct {
        Value       ResultValue
        Reason      string
        HeaderKey   string
        HeaderValue string
}

func (r *SenderIDResult) parse(value ResultValue, params map[string]string) error {
        r.Value = value
        r.Reason = params["reason"]

        for k, v := range params {
                if strings.HasPrefix(k, "header.") {
                        r.HeaderKey = strings.TrimPrefix(k, "header.")
                        r.HeaderValue = v
                        break
                }
        }

        return nil
}

func (r *SenderIDResult) format() (value ResultValue, params map[string]string) {
        return r.Value, map[string]string{
                "reason":                                 r.Reason,
                "header." + strings.ToLower(r.HeaderKey): r.HeaderValue,
        }
}

type SPFResult struct {
        Value  ResultValue
        Reason string
        From   string
        Helo   string
}

func (r *SPFResult) parse(value ResultValue, params map[string]string) error {
        r.Value = value
        r.Reason = params["reason"]
        r.From = params["smtp.mailfrom"]
        r.Helo = params["smtp.helo"]
        return nil
}

func (r *SPFResult) format() (ResultValue, map[string]string) {
        return r.Value, map[string]string{
                "reason":        r.Reason,
                "smtp.mailfrom": r.From,
                "smtp.helo":     r.Helo,
        }
}

type DMARCResult struct {
        Value  ResultValue
        Reason string
        From   string
}

func (r *DMARCResult) parse(value ResultValue, params map[string]string) error {
        r.Value = value
        r.Reason = params["reason"]
        r.From = params["header.from"]
        return nil
}

func (r *DMARCResult) format() (ResultValue, map[string]string) {
        return r.Value, map[string]string{
                "reason":      r.Reason,
                "header.from": r.From,
        }
}

type ARCResult struct {
        Value      ResultValue
        RemoteIP   string
        OldestPass int
}

func (r *ARCResult) parse(value ResultValue, params map[string]string) error {
        var oldestPass int
        if s, ok := params["header.oldest-pass"]; ok {
                var err error
                oldestPass, err = strconv.Atoi(s)
                if err != nil {
                        return fmt.Errorf("invalid header.oldest-pass param: %v", err)
                } else if oldestPass <= 0 {
                        return fmt.Errorf("invalid header.oldest-pass param: must be >= 1")
                }
        }

        r.Value = value
        r.RemoteIP = params["smtp.remote-ip"]
        r.OldestPass = oldestPass
        return nil
}

func (r *ARCResult) format() (ResultValue, map[string]string) {
        var oldestPass string
        if r.OldestPass > 0 {
                oldestPass = strconv.Itoa(r.OldestPass)
        }

        return r.Value, map[string]string{
                "smtp.remote-ip":     r.RemoteIP,
                "header.oldest-pass": oldestPass,
        }
}

type GenericResult struct {
        Method string
        Value  ResultValue
        Params map[string]string
}

func (r *GenericResult) parse(value ResultValue, params map[string]string) error {
        r.Value = value
        r.Params = params
        return nil
}

func (r *GenericResult) format() (ResultValue, map[string]string) {
        return r.Value, r.Params
}

type newResultFunc func() Result

var results = map[string]newResultFunc{
        "arc": func() Result {
                return new(ARCResult)
        },
        "auth": func() Result {
                return new(AuthResult)
        },
        "dkim": func() Result {
                return new(DKIMResult)
        },
        "domainkeys": func() Result {
                return new(DomainKeysResult)
        },
        "iprev": func() Result {
                return new(IPRevResult)
        },
        "sender-id": func() Result {
                return new(SenderIDResult)
        },
        "spf": func() Result {
                return new(SPFResult)
        },
        "dmarc": func() Result {
                return new(DMARCResult)
        },
}

// Parse parses the provided Authentication-Results header field. It returns the
// authentication service identifier and authentication results.
func Parse(v string) (identifier string, results []Result, err error) {
        parts := strings.Split(v, ";")

        identifier = strings.TrimSpace(parts[0])
        i := strings.IndexFunc(identifier, unicode.IsSpace)
        if i > 0 {
                version := strings.TrimSpace(identifier[i:])
                if version != "1" {
                        return "", nil, errors.New("msgauth: unsupported version")
                }

                identifier = identifier[:i]
        }

        for i := 1; i < len(parts); i++ {
                s := strings.TrimSpace(parts[i])
                if s == "" {
                        continue
                }

                result, err := parseResult(s)
                if err != nil {
                        return identifier, results, err
                }
                if result != nil {
                        results = append(results, result)
                }
        }
        return
}

func parseResult(s string) (Result, error) {
        // TODO: ignore header comments in parenthesis

        parts := strings.Fields(s)
        if len(parts) == 0 || parts[0] == "none" {
                return nil, nil
        }

        k, v, err := parseParam(parts[0])
        if err != nil {
                return nil, err
        }
        method, value := k, ResultValue(strings.ToLower(v))

        params := make(map[string]string)
        for i := 1; i < len(parts); i++ {
                k, v, err := parseParam(parts[i])
                if err != nil {
                        continue
                }

                params[k] = v
        }

        newResult, ok := results[method]

        var r Result
        if ok {
                r = newResult()
        } else {
                r = &GenericResult{
                        Method: method,
                        Value:  value,
                        Params: params,
                }
        }

        err = r.parse(value, params)
        return r, err
}

func parseParam(s string) (k string, v string, err error) {
        k, v, ok := strings.Cut(s, "=")
        if !ok {
                return "", "", errors.New("msgauth: malformed authentication method and value")
        }
        return strings.ToLower(strings.TrimSpace(k)), strings.TrimSpace(v), nil
}

package main

import (
        "crypto"
        "crypto/ed25519"
        "crypto/rand"
        "crypto/rsa"
        "crypto/x509"
        "encoding/base64"
        "encoding/pem"
        "flag"
        "fmt"
        "log"
        "os"
        "strings"
)

var (
        keyType  string
        nBits    int
        filename string
        readPriv bool
)

func init() {
        flag.StringVar(&keyType, "t", "rsa", "key type (rsa, ed25519)")
        flag.IntVar(&nBits, "b", 3072, "number of bits in the key (only for RSA)")
        flag.StringVar(&filename, "f", "dkim.priv", "private key filename")
        flag.BoolVar(&readPriv, "y", false, "read private key and print public key")
        flag.Parse()
}

type privateKey interface {
        Public() crypto.PublicKey
}

func main() {
        var privKey privateKey
        if readPriv {
                privKey = readPrivKey()
        } else {
                privKey = genPrivKey()
                writePrivKey(privKey)
        }
        printPubKey(privKey.Public())
}

func genPrivKey() privateKey {
        var (
                privKey crypto.Signer
                err     error
        )
        switch keyType {
        case "rsa":
                log.Printf("Generating a %v-bit RSA key", nBits)
                privKey, err = rsa.GenerateKey(rand.Reader, nBits)
        case "ed25519":
                log.Printf("Generating an Ed25519 key")
                _, privKey, err = ed25519.GenerateKey(rand.Reader)
        default:
                log.Fatalf("Unsupported key type %q", keyType)
        }
        if err != nil {
                log.Fatalf("Failed to generate key: %v", err)
        }
        return privKey
}

func readPrivKey() privateKey {
        b, err := os.ReadFile(filename)
        if err != nil {
                log.Fatalf("Failed to read public key file: %v", err)
        }

        block, _ := pem.Decode(b)
        if block == nil {
                log.Fatalf("Failed to decode PEM block")
        } else if block.Type != "PRIVATE KEY" {
                log.Fatalf("Not a private key")
        }

        privKey, err := x509.ParsePKCS8PrivateKey(block.Bytes)
        if err != nil {
                log.Fatalf("Failed to parse private key: %v", err)
        }

        log.Printf("Private key read from %q", filename)
        return privKey.(privateKey)
}

func writePrivKey(privKey privateKey) {
        privBytes, err := x509.MarshalPKCS8PrivateKey(privKey)
        if err != nil {
                log.Fatalf("Failed to marshal private key: %v", err)
        }

        f, err := os.OpenFile(filename, os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0600)
        if err != nil {
                log.Fatalf("Failed to create key file: %v", err)
        }
        defer f.Close()

        privBlock := pem.Block{
                Type:  "PRIVATE KEY",
                Bytes: privBytes,
        }
        if err := pem.Encode(f, &privBlock); err != nil {
                log.Fatalf("Failed to write key PEM block: %v", err)
        }
        if err := f.Close(); err != nil {
                log.Fatalf("Failed to close key file: %v", err)
        }
        log.Printf("Private key written to %q", filename)
}

func printPubKey(pubKey crypto.PublicKey) {
        var pubBytes []byte
        switch pubKey := pubKey.(type) {
        case *rsa.PublicKey:
                // RFC 6376 is inconsistent about whether RSA public keys should
                // be formatted as RSAPublicKey or SubjectPublicKeyInfo.
                // Erratum 3017 (https://www.rfc-editor.org/errata/eid3017)
                // proposes allowing both.  We use SubjectPublicKeyInfo for
                // consistency with other implementations including opendkim,
                // Gmail, and Fastmail.
                var err error
                pubBytes, err = x509.MarshalPKIXPublicKey(pubKey)
                if err != nil {
                        log.Fatalf("Failed to marshal public key: %v", err)
                }
        case ed25519.PublicKey:
                pubBytes = pubKey
        default:
                panic("unreachable")
        }

        params := []string{
                "v=DKIM1",
                "k=" + keyType,
                "p=" + base64.StdEncoding.EncodeToString(pubBytes),
        }
        log.Println("Public key, to be stored in the TXT record \"<selector>._domainkey\":")
        fmt.Println(strings.Join(params, "; "))
}

package main

import (
        "bytes"
        "crypto"
        "crypto/x509"
        "encoding/pem"
        "flag"
        "fmt"
        "io"
        "io/ioutil"
        "log"
        "net"
        "net/mail"
        "net/textproto"
        "os"
        "os/signal"
        "path"
        "strings"
        "syscall"

        "github.com/emersion/go-milter"
        "github.com/emersion/go-msgauth/authres"
        "github.com/emersion/go-msgauth/dkim"
        "golang.org/x/crypto/ed25519"
)

var (
        signDomains    stringSliceFlag
        identity       string
        listenURI      string
        privateKeyPath string
        selector       string
        verbose        bool
)

var privateKey crypto.Signer

var signHeaderKeys = []string{
        "From",
        "Reply-To",
        "Subject",
        "Date",
        "To",
        "Cc",
        "Resent-Date",
        "Resent-From",
        "Resent-To",
        "Resent-Cc",
        "In-Reply-To",
        "References",
        "List-Id",
        "List-Help",
        "List-Unsubscribe",
        "List-Subscribe",
        "List-Post",
        "List-Owner",
        "List-Archive",
}

const maxVerifications = 5

const usage = `usage: dkim-milter [options....]

Mail filter to verify and sign messages with DKIM.

By default, message signatures are verified and an Authentication-Results
header field is inserted.

When -d, -k and -s are provided, messages matching the domain(s) are signed.
dkim-keygen can be used to generate a private key.

Options:
`

func init() {
        flag.Var(&signDomains, "d", "Domain(s) whose mail should be signed (glob patterns are allowed)")
        flag.StringVar(&identity, "i", "", "Server identity (defaults to hostname)")
        flag.StringVar(&listenURI, "l", "unix:///tmp/dkim-milter.sock", "Listen URI")
        flag.StringVar(&privateKeyPath, "k", "", "Private key (PEM-formatted)")
        flag.StringVar(&selector, "s", "", "Selector")
        flag.BoolVar(&verbose, "v", false, "Enable verbose logging")

        flag.Usage = func() {
                fmt.Fprintf(flag.CommandLine.Output(), usage)
                flag.PrintDefaults()
        }
}

type stringSliceFlag []string

func (f *stringSliceFlag) String() string {
        return strings.Join(*f, ", ")
}

func (f *stringSliceFlag) Set(value string) error {
        *f = append(*f, value)
        return nil
}

type session struct {
        milter.NoOpMilter

        authResDelete []int
        headerBuf     bytes.Buffer

        signDomain     string
        signHeaderKeys []string

        done   <-chan error
        pw     *io.PipeWriter
        verifs []*dkim.Verification // only valid after done is closed
        signer *dkim.Signer
        mw     io.Writer
}

func parseAddressDomain(s string) (string, error) {
        addr, err := mail.ParseAddress(s)
        if err != nil {
                return "", err
        }

        parts := strings.SplitN(addr.Address, "@", 2)
        if len(parts) != 2 {
                return "", fmt.Errorf("dkim-milter: malformed address: missing '@'")
        }

        return parts[1], nil
}

func (s *session) Header(name string, value string, m *milter.Modifier) (milter.Response, error) {
        if strings.EqualFold(name, "From") || strings.EqualFold(name, "Sender") {
                domain, err := parseAddressDomain(value)
                if err != nil {
                        return nil, fmt.Errorf("dkim-milter: failed to parse header field %q: %v", name, err)
                }
                domain = strings.ToLower(domain)

                for _, pattern := range signDomains {
                        if ok, err := path.Match(pattern, domain); err != nil {
                                return nil, fmt.Errorf("dkim-milter: failed to match domain %q: %v", domain, err)
                        } else if ok {
                                s.signDomain = domain
                                break
                        }
                }
        }

        for _, k := range signHeaderKeys {
                if strings.EqualFold(name, k) {
                        s.signHeaderKeys = append(s.signHeaderKeys, name)
                }
        }

        field := name + ": " + value + "\r\n"
        _, err := s.headerBuf.WriteString(field)
        return milter.RespContinue, err
}

func getIdentity(authRes string) string {
        parts := strings.SplitN(authRes, ";", 2)
        return strings.TrimSpace(parts[0])
}

func shouldDeleteAuthRes(field string) bool {
        id, results, err := authres.Parse(field)
        if err != nil {
                // Delete fields we can't parse, because other implementations might
                // accept malformed fields
                return true
        }

        if !strings.EqualFold(id, identity) {
                // Not our Authentication-Results, ignore the field
                return false
        }

        for _, res := range results {
                if _, ok := res.(*authres.DKIMResult); ok {
                        // Delete existing DKIM Authentication-Results fields
                        return true
                }
        }

        // This is our Authentication-Results field, but it isn't about DKIM. Maybe
        // a previous milter has generated it (e.g. SPF), so keep it.
        return false
}

func (s *session) Headers(h textproto.MIMEHeader, m *milter.Modifier) (milter.Response, error) {
        // Write final CRLF to begin message body
        if _, err := s.headerBuf.WriteString("\r\n"); err != nil {
                return nil, err
        }

        // Delete any existing Authentication-Results header field with our identity
        fields := h["Authentication-Results"]
        for i, field := range fields {
                if shouldDeleteAuthRes(field) {
                        s.authResDelete = append(s.authResDelete, i)
                }
        }

        // Sign if necessary
        if s.signDomain != "" {
                opts := dkim.SignOptions{
                        Domain:       s.signDomain,
                        Selector:     selector,
                        Signer:       privateKey,
                        HeaderKeys:   s.signHeaderKeys,
                        QueryMethods: []dkim.QueryMethod{dkim.QueryMethodDNSTXT},
                }

                var err error
                s.signer, err = dkim.NewSigner(&opts)
                if err != nil {
                        return nil, err
                }
        }

        // Verify existing signatures
        done := make(chan error, 1)
        pr, pw := io.Pipe()

        s.done = done
        s.pw = pw

        // TODO: limit max. number of signatures
        go func() {
                options := dkim.VerifyOptions{MaxVerifications: maxVerifications}

                var err error
                s.verifs, err = dkim.VerifyWithOptions(pr, &options)
                io.Copy(ioutil.Discard, pr)
                pr.Close()
                done <- err
                close(done)
        }()

        // Process header
        return s.BodyChunk(s.headerBuf.Bytes(), m)
}

func (s *session) BodyChunk(chunk []byte, m *milter.Modifier) (milter.Response, error) {
        if _, err := s.pw.Write(chunk); err != nil {
                return nil, err
        }
        if s.signer != nil {
                if _, err := s.signer.Write(chunk); err != nil {
                        return nil, err
                }
        }
        return milter.RespContinue, nil
}

func (s *session) Body(m *milter.Modifier) (milter.Response, error) {
        if err := s.pw.Close(); err != nil {
                return nil, err
        }

        for _, index := range s.authResDelete {
                if err := m.ChangeHeader(index, "Authentication-Results", ""); err != nil {
                        return nil, err
                }
        }

        if err := <-s.done; err == dkim.ErrTooManySignatures {
                if verbose {
                        log.Printf("Too many signatures in message: %v", err)
                }
                // Ignore the error
        } else if err != nil {
                if verbose {
                        log.Printf("DKIM verification failed: %v", err)
                }
                return nil, err
        }

        if s.signer != nil {
                if err := s.signer.Close(); err != nil {
                        if verbose {
                                log.Printf("DKIM signature failed: %v", err)
                        }
                        return nil, err
                }

                kv := s.signer.Signature()
                parts := strings.SplitN(kv, ": ", 2)
                if len(parts) != 2 {
                        panic("dkim-milter: malformed DKIM-Signature header field")
                }
                k, v := parts[0], strings.TrimSuffix(parts[1], "\r\n")

                if err := m.InsertHeader(0, k, v); err != nil {
                        return nil, err
                }
        }

        results := make([]authres.Result, 0, len(s.verifs))

        if len(s.verifs) == 0 && s.signer == nil {
                results = append(results, &authres.DKIMResult{
                        Value: authres.ResultNone,
                })
        }

        for _, verif := range s.verifs {
                if verbose {
                        if verif.Err != nil {
                                log.Printf("DKIM verification failed for %v: %v", verif.Domain, verif.Err)
                        } else {
                                log.Printf("DKIM verification succeded for %v", verif.Domain)
                        }
                }

                var val authres.ResultValue
                if verif.Err == nil {
                        val = authres.ResultPass
                } else if dkim.IsPermFail(verif.Err) {
                        val = authres.ResultPermError
                } else if dkim.IsTempFail(verif.Err) {
                        val = authres.ResultTempError
                } else {
                        val = authres.ResultFail
                }

                results = append(results, &authres.DKIMResult{
                        Value:      val,
                        Domain:     verif.Domain,
                        Identifier: verif.Identifier,
                })
        }

        if len(s.verifs) > 0 || s.signer == nil {
                v := authres.Format(identity, results)
                if err := m.InsertHeader(0, "Authentication-Results", v); err != nil {
                        return nil, err
                }
        }

        return milter.RespAccept, nil
}

func loadPrivateKey(path string) (crypto.Signer, error) {
        b, err := ioutil.ReadFile(privateKeyPath)
        if err != nil {
                return nil, err
        }

        block, _ := pem.Decode(b)
        if block == nil {
                return nil, fmt.Errorf("no PEM data found")
        }

        switch strings.ToUpper(block.Type) {
        case "PRIVATE KEY":
                k, err := x509.ParsePKCS8PrivateKey(block.Bytes)
                if err != nil {
                        return nil, err
                }
                return k.(crypto.Signer), nil
        case "RSA PRIVATE KEY":
                return x509.ParsePKCS1PrivateKey(block.Bytes)
        case "EDDSA PRIVATE KEY":
                if len(block.Bytes) != ed25519.PrivateKeySize {
                        return nil, fmt.Errorf("invalid Ed25519 private key size")
                }
                return ed25519.PrivateKey(block.Bytes), nil
        default:
                return nil, fmt.Errorf("unknown private key type: '%v'", block.Type)
        }
}

func main() {
        flag.Parse()

        if identity == "" {
                var err error
                identity, err = os.Hostname()
                if err != nil {
                        log.Fatal("Failed to read hostname: ", err)
                }
        }

        if (len(signDomains) > 0 || privateKeyPath != "" || selector != "") && !(len(signDomains) > 0 && privateKeyPath != "" && selector != "") {
                log.Fatal("Domain(s) (-d), private key (-k) and selector (-s) must all be specified")
        }

        for i, pattern := range signDomains {
                if _, err := path.Match(pattern, ""); err != nil {
                        log.Fatalf("Malformed domain pattern %q: %v", pattern, err)
                }
                signDomains[i] = strings.ToLower(pattern)
        }

        if privateKeyPath != "" {
                var err error
                privateKey, err = loadPrivateKey(privateKeyPath)
                if err != nil {
                        log.Fatalf("Failed to load private key from '%v': %v", privateKeyPath, err)
                }
        }

        parts := strings.SplitN(listenURI, "://", 2)
        if len(parts) != 2 {
                log.Fatal("Invalid listen URI")
        }
        listenNetwork, listenAddr := parts[0], parts[1]

        s := milter.Server{
                NewMilter: func() milter.Milter {
                        return &session{}
                },
                Actions:  milter.OptAddHeader | milter.OptChangeHeader,
                Protocol: milter.OptNoConnect | milter.OptNoHelo | milter.OptNoMailFrom | milter.OptNoRcptTo,
        }

        ln, err := net.Listen(listenNetwork, listenAddr)
        if err != nil {
                log.Fatal("Failed to setup listener: ", err)
        }

        // Closing the listener will unlink the unix socket, if any
        sigs := make(chan os.Signal, 1)
        signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)
        go func() {
                <-sigs
                if err := s.Close(); err != nil {
                        log.Fatal("Failed to close server: ", err)
                }
        }()

        log.Println("Milter listening at", listenURI)
        if err := s.Serve(ln); err != nil && err != milter.ErrServerClosed {
                log.Fatal("Failed to serve: ", err)
        }
}

package main

import (
        "log"
        "os"

        "github.com/emersion/go-msgauth/dkim"
)

func main() {
        verifications, err := dkim.Verify(os.Stdin)
        if err != nil {
                log.Fatal(err)
        }

        for _, v := range verifications {
                if v.Err == nil {
                        log.Printf("Valid signature for %v", v.Domain)
                } else {
                        log.Printf("Invalid signature for %v: %v", v.Domain, v.Err)
                }
        }
}

package main

import (
        "flag"
        "log"

        "github.com/emersion/go-msgauth/dmarc"
)

func main() {
        flag.Parse()

        domain := flag.Arg(0)
        if domain == "" {
                log.Fatal("usage: dmarc-lookup <domain>")
        }

        rec, err := dmarc.Lookup(domain)
        if err != nil {
                log.Fatal(err)
        }

        log.Printf("%#v\n", rec)
}

package dkim

import (
        "io"
        "strings"
)

// Canonicalization is a canonicalization algorithm.
type Canonicalization string

const (
        CanonicalizationSimple  Canonicalization = "simple"
        CanonicalizationRelaxed                  = "relaxed"
)

type canonicalizer interface {
        CanonicalizeHeader(s string) string
        CanonicalizeBody(w io.Writer) io.WriteCloser
}

var canonicalizers = map[Canonicalization]canonicalizer{
        CanonicalizationSimple:  new(simpleCanonicalizer),
        CanonicalizationRelaxed: new(relaxedCanonicalizer),
}

// crlfFixer fixes any lone LF without a preceding CR.
type crlfFixer struct {
        cr bool
}

func (cf *crlfFixer) Fix(b []byte) []byte {
        res := make([]byte, 0, len(b))
        for _, ch := range b {
                prevCR := cf.cr
                cf.cr = false
                switch ch {
                case '\r':
                        cf.cr = true
                case '\n':
                        if !prevCR {
                                res = append(res, '\r')
                        }
                }
                res = append(res, ch)
        }
        return res
}

type simpleCanonicalizer struct{}

func (c *simpleCanonicalizer) CanonicalizeHeader(s string) string {
        return s
}

type simpleBodyCanonicalizer struct {
        w         io.Writer
        crlfBuf   []byte
        crlfFixer crlfFixer
}

func (c *simpleBodyCanonicalizer) Write(b []byte) (int, error) {
        written := len(b)
        b = append(c.crlfBuf, b...)

        b = c.crlfFixer.Fix(b)

        end := len(b)
        // If it ends with \r, maybe the next write will begin with \n
        if end > 0 && b[end-1] == '\r' {
                end--
        }
        // Keep all \r\n sequences
        for end >= 2 {
                prev := b[end-2]
                cur := b[end-1]
                if prev != '\r' || cur != '\n' {
                        break
                }
                end -= 2
        }

        c.crlfBuf = b[end:]

        var err error
        if end > 0 {
                _, err = c.w.Write(b[:end])
        }
        return written, err
}

func (c *simpleBodyCanonicalizer) Close() error {
        // Flush crlfBuf if it ends with a single \r (without a matching \n)
        if len(c.crlfBuf) > 0 && c.crlfBuf[len(c.crlfBuf)-1] == '\r' {
                if _, err := c.w.Write(c.crlfBuf); err != nil {
                        return err
                }
        }
        c.crlfBuf = nil

        if _, err := c.w.Write([]byte(crlf)); err != nil {
                return err
        }
        return nil
}

func (c *simpleCanonicalizer) CanonicalizeBody(w io.Writer) io.WriteCloser {
        return &simpleBodyCanonicalizer{w: w}
}

type relaxedCanonicalizer struct{}

func (c *relaxedCanonicalizer) CanonicalizeHeader(s string) string {
        k, v, ok := strings.Cut(s, ":")
        if !ok {
                return strings.TrimSpace(strings.ToLower(s)) + ":" + crlf
        }

        k = strings.TrimSpace(strings.ToLower(k))
        v = strings.Join(strings.FieldsFunc(v, func(r rune) bool {
                return r == ' ' || r == '\t' || r == '\n' || r == '\r'
        }), " ")
        return k + ":" + v + crlf
}

type relaxedBodyCanonicalizer struct {
        w         io.Writer
        crlfBuf   []byte
        wsp       bool
        written   bool
        crlfFixer crlfFixer
}

func (c *relaxedBodyCanonicalizer) Write(b []byte) (int, error) {
        written := len(b)

        b = c.crlfFixer.Fix(b)

        canonical := make([]byte, 0, len(b))
        for _, ch := range b {
                if ch == ' ' || ch == '\t' {
                        c.wsp = true
                } else if ch == '\r' || ch == '\n' {
                        c.wsp = false
                        c.crlfBuf = append(c.crlfBuf, ch)
                } else {
                        if len(c.crlfBuf) > 0 {
                                canonical = append(canonical, c.crlfBuf...)
                                c.crlfBuf = c.crlfBuf[:0]
                        }
                        if c.wsp {
                                canonical = append(canonical, ' ')
                                c.wsp = false
                        }

                        canonical = append(canonical, ch)
                }
        }

        if !c.written && len(canonical) > 0 {
                c.written = true
        }

        _, err := c.w.Write(canonical)
        return written, err
}

func (c *relaxedBodyCanonicalizer) Close() error {
        if c.written {
                if _, err := c.w.Write([]byte(crlf)); err != nil {
                        return err
                }
        }
        return nil
}

func (c *relaxedCanonicalizer) CanonicalizeBody(w io.Writer) io.WriteCloser {
        return &relaxedBodyCanonicalizer{w: w}
}

type limitedWriter struct {
        W io.Writer
        N int64
}

func (w *limitedWriter) Write(b []byte) (int, error) {
        if w.N <= 0 {
                return len(b), nil
        }

        skipped := 0
        if int64(len(b)) > w.N {
                b = b[:w.N]
                skipped = int(int64(len(b)) - w.N)
        }

        n, err := w.W.Write(b)
        w.N -= int64(n)
        return n + skipped, err
}

package dkim

import (
        "bufio"
        "bytes"
        "errors"
        "fmt"
        "io"
        "net/textproto"
        "sort"
        "strings"
)

const crlf = "\r\n"

type header []string

func readHeader(r *bufio.Reader) (header, error) {
        tr := textproto.NewReader(r)

        var h header
        for {
                l, err := tr.ReadLine()
                if err != nil {
                        return h, fmt.Errorf("failed to read header: %v", err)
                }

                if len(l) == 0 {
                        break
                } else if len(h) > 0 && (l[0] == ' ' || l[0] == '\t') {
                        // This is a continuation line
                        h[len(h)-1] += l + crlf
                } else {
                        h = append(h, l+crlf)
                }
        }

        return h, nil
}

func writeHeader(w io.Writer, h header) error {
        for _, kv := range h {
                if _, err := w.Write([]byte(kv)); err != nil {
                        return err
                }
        }
        _, err := w.Write([]byte(crlf))
        return err
}

func foldHeaderField(kv string) string {
        buf := bytes.NewBufferString(kv)

        line := make([]byte, 75) // 78 - len("\r\n\s")
        first := true
        var fold strings.Builder
        for len, err := buf.Read(line); err != io.EOF; len, err = buf.Read(line) {
                if first {
                        first = false
                } else {
                        fold.WriteString("\r\n ")
                }
                fold.Write(line[:len])
        }

        return fold.String() + crlf
}

func parseHeaderField(s string) (string, string) {
        key, value, _ := strings.Cut(s, ":")
        return strings.TrimSpace(key), strings.TrimSpace(value)
}

func parseHeaderParams(s string) (map[string]string, error) {
        pairs := strings.Split(s, ";")
        params := make(map[string]string)
        for _, s := range pairs {
                key, value, ok := strings.Cut(s, "=")
                if !ok {
                        if strings.TrimSpace(s) == "" {
                                continue
                        }
                        return params, errors.New("dkim: malformed header params")
                }

                trimmedKey := strings.TrimSpace(key)
                _, present := params[trimmedKey]
                if present {
                        return params, errors.New("dkim: duplicate tag name")
                }
                params[trimmedKey] = strings.TrimSpace(value)
        }
        return params, nil
}

func formatHeaderParams(headerFieldName string, params map[string]string) string {
        keys, bvalue, bfound := sortParams(params)

        s := headerFieldName + ":"
        var line string

        for _, k := range keys {
                v := params[k]
                nextLength := 3 + len(line) + len(v) + len(k)
                if nextLength > 75 {
                        s += line + crlf
                        line = ""
                }
                line = fmt.Sprintf("%v %v=%v;", line, k, v)
        }

        if line != "" {
                s += line
        }

        if bfound {
                bfiled := foldHeaderField(" b=" + bvalue)
                s += crlf + bfiled
        }

        return s
}

func sortParams(params map[string]string) ([]string, string, bool) {
        keys := make([]string, 0, len(params))
        bfound := false
        var bvalue string
        for k := range params {
                if k == "b" {
                        bvalue = params["b"]
                        bfound = true
                } else {
                        keys = append(keys, k)
                }
        }
        sort.Strings(keys)
        return keys, bvalue, bfound
}

type headerPicker struct {
        h      header
        picked map[string]int
}

func newHeaderPicker(h header) *headerPicker {
        return &headerPicker{
                h:      h,
                picked: make(map[string]int),
        }
}

func (p *headerPicker) Pick(key string) string {
        key = strings.ToLower(key)

        at := p.picked[key]
        for i := len(p.h) - 1; i >= 0; i-- {
                kv := p.h[i]
                k, _ := parseHeaderField(kv)

                if !strings.EqualFold(k, key) {
                        continue
                }

                if at == 0 {
                        p.picked[key]++
                        return kv
                }
                at--
        }

        return ""
}

package dkim

import (
        "crypto"
        "crypto/rsa"
        "crypto/x509"
        "encoding/base64"
        "errors"
        "fmt"
        "net"
        "strings"

        "golang.org/x/crypto/ed25519"
)

type verifier interface {
        Public() crypto.PublicKey
        Verify(hash crypto.Hash, hashed []byte, sig []byte) error
}

type rsaVerifier struct {
        *rsa.PublicKey
}

func (v rsaVerifier) Public() crypto.PublicKey {
        return v.PublicKey
}

func (v rsaVerifier) Verify(hash crypto.Hash, hashed, sig []byte) error {
        return rsa.VerifyPKCS1v15(v.PublicKey, hash, hashed, sig)
}

type ed25519Verifier struct {
        ed25519.PublicKey
}

func (v ed25519Verifier) Public() crypto.PublicKey {
        return v.PublicKey
}

func (v ed25519Verifier) Verify(hash crypto.Hash, hashed, sig []byte) error {
        if !ed25519.Verify(v.PublicKey, hashed, sig) {
                return errors.New("dkim: invalid Ed25519 signature")
        }
        return nil
}

type queryResult struct {
        Verifier  verifier
        KeyAlgo   string
        HashAlgos []string
        Notes     string
        Services  []string
        Flags     []string
}

// QueryMethod is a DKIM query method.
type QueryMethod string

const (
        // DNS TXT resource record (RR) lookup algorithm
        QueryMethodDNSTXT QueryMethod = "dns/txt"
)

type txtLookupFunc func(domain string) ([]string, error)
type queryFunc func(domain, selector string, txtLookup txtLookupFunc) (*queryResult, error)

var queryMethods = map[QueryMethod]queryFunc{
        QueryMethodDNSTXT: queryDNSTXT,
}

func queryDNSTXT(domain, selector string, txtLookup txtLookupFunc) (*queryResult, error) {
        if txtLookup == nil {
                txtLookup = net.LookupTXT
        }

        txts, err := txtLookup(selector + "._domainkey." + domain)
        if netErr, ok := err.(net.Error); ok && netErr.Temporary() {
                return nil, tempFailError("key unavailable: " + err.Error())
        } else if err != nil {
                return nil, permFailError("no key for signature: " + err.Error())
        }

        // net.LookupTXT will concatenate strings contained in a single TXT record.
        // In other words, net.LookupTXT returns one entry per TXT record, even if
        // a record contains multiple strings.
        //
        // RFC 6376 section 3.6.2.2 says multiple TXT records lead to undefined
        // behavior, so reject that.
        switch len(txts) {
        case 0:
                return nil, permFailError("no valid key found")
        case 1:
                return parsePublicKey(txts[0])
        default:
                return nil, permFailError("multiple TXT records found for key")
        }
}

func parsePublicKey(s string) (*queryResult, error) {
        params, err := parseHeaderParams(s)
        if err != nil {
                return nil, permFailError("key record error: " + err.Error())
        }

        res := new(queryResult)

        if v, ok := params["v"]; ok && v != "DKIM1" {
                return nil, permFailError("incompatible public key version")
        }

        p, ok := params["p"]
        if !ok {
                return nil, permFailError("key syntax error: missing public key data")
        }
        if p == "" {
                return nil, permFailError("key revoked")
        }
        p = strings.ReplaceAll(p, " ", "")
        b, err := base64.StdEncoding.DecodeString(p)
        if err != nil {
                return nil, permFailError("key syntax error: " + err.Error())
        }
        switch params["k"] {
        case "rsa", "":
                pub, err := x509.ParsePKIXPublicKey(b)
                if err != nil {
                        // RFC 6376 is inconsistent about whether RSA public keys should
                        // be formatted as RSAPublicKey or SubjectPublicKeyInfo.
                        // Erratum 3017 (https://www.rfc-editor.org/errata/eid3017) proposes
                        // allowing both.
                        pub, err = x509.ParsePKCS1PublicKey(b)
                        if err != nil {
                                return nil, permFailError("key syntax error: " + err.Error())
                        }
                }
                rsaPub, ok := pub.(*rsa.PublicKey)
                if !ok {
                        return nil, permFailError("key syntax error: not an RSA public key")
                }
                // RFC 8301 section 3.2: verifiers MUST NOT consider signatures using
                // RSA keys of less than 1024 bits as valid signatures.
                if rsaPub.Size()*8 < 1024 {
                        return nil, permFailError(fmt.Sprintf("key is too short: want 1024 bits, has %v bits", rsaPub.Size()*8))
                }
                res.Verifier = rsaVerifier{rsaPub}
                res.KeyAlgo = "rsa"
        case "ed25519":
                if len(b) != ed25519.PublicKeySize {
                        return nil, permFailError(fmt.Sprintf("invalid Ed25519 public key size: %v bytes", len(b)))
                }
                ed25519Pub := ed25519.PublicKey(b)
                res.Verifier = ed25519Verifier{ed25519Pub}
                res.KeyAlgo = "ed25519"
        default:
                return nil, permFailError("unsupported key algorithm")
        }

        if hashesStr, ok := params["h"]; ok {
                res.HashAlgos = parseTagList(hashesStr)
        }
        if notes, ok := params["n"]; ok {
                res.Notes = notes
        }
        if servicesStr, ok := params["s"]; ok {
                services := parseTagList(servicesStr)

                hasWildcard := false
                for _, s := range services {
                        if s == "*" {
                                hasWildcard = true
                                break
                        }
                }
                if !hasWildcard {
                        res.Services = services
                }
        }
        if flagsStr, ok := params["t"]; ok {
                res.Flags = parseTagList(flagsStr)
        }

        return res, nil
}

package dkim

import (
        "bufio"
        "bytes"
        "crypto"
        "crypto/rand"
        "crypto/rsa"
        "encoding/base64"
        "fmt"
        "io"
        "strconv"
        "strings"
        "time"

        "golang.org/x/crypto/ed25519"
)

var randReader io.Reader = rand.Reader

// SignOptions is used to configure Sign. Domain, Selector and Signer are
// mandatory.
type SignOptions struct {
        // The SDID claiming responsibility for an introduction of a message into the
        // mail stream. Hence, the SDID value is used to form the query for the public
        // key. The SDID MUST correspond to a valid DNS name under which the DKIM key
        // record is published.
        //
        // This can't be empty.
        Domain string
        // The selector subdividing the namespace for the domain.
        //
        // This can't be empty.
        Selector string
        // The Agent or User Identifier (AUID) on behalf of which the SDID is taking
        // responsibility.
        //
        // This is optional.
        Identifier string

        // The key used to sign the message.
        //
        // Supported Signer.Public() values are *rsa.PublicKey and
        // ed25519.PublicKey.
        Signer crypto.Signer
        // The hash algorithm used to sign the message. If zero, a default hash will
        // be chosen.
        //
        // The only supported hash algorithm is crypto.SHA256.
        Hash crypto.Hash

        // Header and body canonicalization algorithms.
        //
        // If empty, CanonicalizationSimple is used.
        HeaderCanonicalization Canonicalization
        BodyCanonicalization   Canonicalization

        // A list of header fields to include in the signature. If nil, all headers
        // will be included. If not nil, "From" MUST be in the list.
        //
        // See RFC 6376 section 5.4.1 for recommended header fields.
        HeaderKeys []string

        // The expiration time. A zero value means no expiration.
        Expiration time.Time

        // A list of query methods used to retrieve the public key.
        //
        // If nil, it is implicitly defined as QueryMethodDNSTXT.
        QueryMethods []QueryMethod
}

// Signer generates a DKIM signature.
//
// The whole message header and body must be written to the Signer. Close should
// always be called (either after the whole message has been written, or after
// an error occurred and the signer won't be used anymore). Close may return an
// error in case signing fails.
//
// After a successful Close, Signature can be called to retrieve the
// DKIM-Signature header field that the caller should prepend to the message.
type Signer struct {
        pw        *io.PipeWriter
        done      <-chan error
        sigParams map[string]string // only valid after done received nil
}

// NewSigner creates a new signer. It returns an error if SignOptions is
// invalid.
func NewSigner(options *SignOptions) (*Signer, error) {
        if options == nil {
                return nil, fmt.Errorf("dkim: no options specified")
        }
        if options.Domain == "" {
                return nil, fmt.Errorf("dkim: no domain specified")
        }
        if options.Selector == "" {
                return nil, fmt.Errorf("dkim: no selector specified")
        }
        if options.Signer == nil {
                return nil, fmt.Errorf("dkim: no signer specified")
        }

        headerCan := options.HeaderCanonicalization
        if headerCan == "" {
                headerCan = CanonicalizationSimple
        }
        if _, ok := canonicalizers[headerCan]; !ok {
                return nil, fmt.Errorf("dkim: unknown header canonicalization %q", headerCan)
        }

        bodyCan := options.BodyCanonicalization
        if bodyCan == "" {
                bodyCan = CanonicalizationSimple
        }
        if _, ok := canonicalizers[bodyCan]; !ok {
                return nil, fmt.Errorf("dkim: unknown body canonicalization %q", bodyCan)
        }

        var keyAlgo string
        switch options.Signer.Public().(type) {
        case *rsa.PublicKey:
                keyAlgo = "rsa"
        case ed25519.PublicKey:
                keyAlgo = "ed25519"
        default:
                return nil, fmt.Errorf("dkim: unsupported key algorithm %T", options.Signer.Public())
        }

        hash := options.Hash
        var hashAlgo string
        switch options.Hash {
        case 0: // sha256 is the default
                hash = crypto.SHA256
                fallthrough
        case crypto.SHA256:
                hashAlgo = "sha256"
        case crypto.SHA1:
                return nil, fmt.Errorf("dkim: hash algorithm too weak: sha1")
        default:
                return nil, fmt.Errorf("dkim: unsupported hash algorithm")
        }

        if options.HeaderKeys != nil {
                ok := false
                for _, k := range options.HeaderKeys {
                        if strings.EqualFold(k, "From") {
                                ok = true
                                break
                        }
                }
                if !ok {
                        return nil, fmt.Errorf("dkim: the From header field must be signed")
                }
        }

        done := make(chan error, 1)
        pr, pw := io.Pipe()

        s := &Signer{
                pw:   pw,
                done: done,
        }

        closeReadWithError := func(err error) {
                pr.CloseWithError(err)
                done <- err
        }

        go func() {
                defer close(done)

                // Read header
                br := bufio.NewReader(pr)
                h, err := readHeader(br)
                if err != nil {
                        closeReadWithError(err)
                        return
                }

                // Hash body
                hasher := hash.New()
                can := canonicalizers[bodyCan].CanonicalizeBody(hasher)
                if _, err := io.Copy(can, br); err != nil {
                        closeReadWithError(err)
                        return
                }
                if err := can.Close(); err != nil {
                        closeReadWithError(err)
                        return
                }
                bodyHashed := hasher.Sum(nil)

                params := map[string]string{
                        "v":  "1",
                        "a":  keyAlgo + "-" + hashAlgo,
                        "bh": base64.StdEncoding.EncodeToString(bodyHashed),
                        "c":  string(headerCan) + "/" + string(bodyCan),
                        "d":  options.Domain,
                        //"l": "", // TODO
                        "s": options.Selector,
                        "t": formatTime(now()),
                        //"z": "", // TODO
                }

                var headerKeys []string
                if options.HeaderKeys != nil {
                        headerKeys = options.HeaderKeys
                } else {
                        for _, kv := range h {
                                k, _ := parseHeaderField(kv)
                                headerKeys = append(headerKeys, k)
                        }
                }
                params["h"] = formatTagList(headerKeys)

                if options.Identifier != "" {
                        params["i"] = options.Identifier
                }

                if options.QueryMethods != nil {
                        methods := make([]string, len(options.QueryMethods))
                        for i, method := range options.QueryMethods {
                                methods[i] = string(method)
                        }
                        params["q"] = formatTagList(methods)
                }

                if !options.Expiration.IsZero() {
                        params["x"] = formatTime(options.Expiration)
                }

                // Hash and sign headers
                hasher.Reset()
                picker := newHeaderPicker(h)
                for _, k := range headerKeys {
                        kv := picker.Pick(k)
                        if kv == "" {
                                // The Signer MAY include more instances of a header field name
                                // in "h=" than there are actual corresponding header fields so
                                // that the signature will not verify if additional header
                                // fields of that name are added.
                                continue
                        }

                        kv = canonicalizers[headerCan].CanonicalizeHeader(kv)
                        if _, err := io.WriteString(hasher, kv); err != nil {
                                closeReadWithError(err)
                                return
                        }
                }

                params["b"] = ""
                sigField := formatSignature(params)
                sigField = canonicalizers[headerCan].CanonicalizeHeader(sigField)
                sigField = strings.TrimRight(sigField, crlf)
                if _, err := io.WriteString(hasher, sigField); err != nil {
                        closeReadWithError(err)
                        return
                }
                hashed := hasher.Sum(nil)

                // Don't pass Hash to Sign for ed25519 as it doesn't support it
                // and will return an error ("ed25519: cannot sign hashed message").
                if keyAlgo == "ed25519" {
                        hash = crypto.Hash(0)
                }

                sig, err := options.Signer.Sign(randReader, hashed, hash)
                if err != nil {
                        closeReadWithError(err)
                        return
                }
                params["b"] = base64.StdEncoding.EncodeToString(sig)

                s.sigParams = params
                closeReadWithError(nil)
        }()

        return s, nil
}

// Write implements io.WriteCloser.
func (s *Signer) Write(b []byte) (n int, err error) {
        return s.pw.Write(b)
}

// Close implements io.WriteCloser. The error return by Close must be checked.
func (s *Signer) Close() error {
        if err := s.pw.Close(); err != nil {
                return err
        }
        return <-s.done
}

// Signature returns the whole DKIM-Signature header field. It can only be
// called after a successful Signer.Close call.
//
// The returned value contains both the header field name, its value and the
// final CRLF.
func (s *Signer) Signature() string {
        if s.sigParams == nil {
                panic("dkim: Signer.Signature must only be called after a succesful Signer.Close")
        }
        return formatSignature(s.sigParams)
}

// Sign signs a message. It reads it from r and writes the signed version to w.
func Sign(w io.Writer, r io.Reader, options *SignOptions) error {
        s, err := NewSigner(options)
        if err != nil {
                return err
        }
        defer s.Close()

        // We need to keep the message in a buffer so we can write the new DKIM
        // header field before the rest of the message
        var b bytes.Buffer
        mw := io.MultiWriter(&b, s)

        if _, err := io.Copy(mw, r); err != nil {
                return err
        }
        if err := s.Close(); err != nil {
                return err
        }

        if _, err := io.WriteString(w, s.Signature()); err != nil {
                return err
        }
        _, err = io.Copy(w, &b)
        return err
}

func formatSignature(params map[string]string) string {
        sig := formatHeaderParams(headerFieldName, params)
        return sig
}

func formatTagList(l []string) string {
        return strings.Join(l, ":")
}

func formatTime(t time.Time) string {
        return strconv.FormatInt(t.Unix(), 10)
}

package dkim

import (
        "bufio"
        "crypto"
        "crypto/subtle"
        "encoding/base64"
        "errors"
        "fmt"
        "io"
        "io/ioutil"
        "regexp"
        "strconv"
        "strings"
        "time"
        "unicode"
)

type permFailError string

func (err permFailError) Error() string {
        return "dkim: " + string(err)
}

// IsPermFail returns true if the error returned by Verify is a permanent
// failure. A permanent failure is for instance a missing required field or a
// malformed header.
func IsPermFail(err error) bool {
        _, ok := err.(permFailError)
        return ok
}

type tempFailError string

func (err tempFailError) Error() string {
        return "dkim: " + string(err)
}

// IsTempFail returns true if the error returned by Verify is a temporary
// failure.
func IsTempFail(err error) bool {
        _, ok := err.(tempFailError)
        return ok
}

type failError string

func (err failError) Error() string {
        return "dkim: " + string(err)
}

// isFail returns true if the error returned by Verify is a signature error.
func isFail(err error) bool {
        _, ok := err.(failError)
        return ok
}

// ErrTooManySignatures is returned by Verify when the message exceeds the
// maximum number of signatures.
var ErrTooManySignatures = errors.New("dkim: too many signatures")

var requiredTags = []string{"v", "a", "b", "bh", "d", "h", "s"}

// A Verification is produced by Verify when it checks if one signature is
// valid. If the signature is valid, Err is nil.
type Verification struct {
        // The SDID claiming responsibility for an introduction of a message into the
        // mail stream.
        Domain string
        // The Agent or User Identifier (AUID) on behalf of which the SDID is taking
        // responsibility.
        Identifier string

        // The list of signed header fields.
        HeaderKeys []string

        // The time that this signature was created. If unknown, it's set to zero.
        Time time.Time
        // The expiration time. If the signature doesn't expire, it's set to zero.
        Expiration time.Time

        // Err is nil if the signature is valid.
        Err error
}

type signature struct {
        i int
        v string
}

// VerifyOptions allows to customize the default signature verification
// behavior.
type VerifyOptions struct {
        // LookupTXT returns the DNS TXT records for the given domain name. If nil,
        // net.LookupTXT is used.
        LookupTXT func(domain string) ([]string, error)
        // MaxVerifications controls the maximum number of signature verifications
        // to perform. If more signatures are present, the first MaxVerifications
        // signatures are verified, the rest are ignored and ErrTooManySignatures
        // is returned. If zero, there is no maximum.
        MaxVerifications int
}

// Verify checks if a message's signatures are valid. It returns one
// verification per signature.
//
// There is no guarantee that the reader will be completely consumed.
func Verify(r io.Reader) ([]*Verification, error) {
        return VerifyWithOptions(r, nil)
}

// VerifyWithOptions performs the same task as Verify, but allows specifying
// verification options.
func VerifyWithOptions(r io.Reader, options *VerifyOptions) ([]*Verification, error) {
        // Read header
        bufr := bufio.NewReader(r)
        h, err := readHeader(bufr)
        if err != nil {
                return nil, err
        }

        // Scan header fields for signatures
        var signatures []*signature
        for i, kv := range h {
                k, v := parseHeaderField(kv)
                if strings.EqualFold(k, headerFieldName) {
                        signatures = append(signatures, &signature{i, v})
                }
        }

        tooManySignatures := false
        if options != nil && options.MaxVerifications > 0 && len(signatures) > options.MaxVerifications {
                tooManySignatures = true
                signatures = signatures[:options.MaxVerifications]
        }

        var verifs []*Verification
        if len(signatures) == 1 {
                // If there is only one signature - just verify it.
                v, err := verify(h, bufr, h[signatures[0].i], signatures[0].v, options)
                if err != nil && !IsTempFail(err) && !IsPermFail(err) && !isFail(err) {
                        return nil, err
                }
                v.Err = err
                verifs = []*Verification{v}
        } else {
                verifs, err = parallelVerify(bufr, h, signatures, options)
                if err != nil {
                        return nil, err
                }
        }

        if tooManySignatures {
                return verifs, ErrTooManySignatures
        }
        return verifs, nil
}

func parallelVerify(r io.Reader, h header, signatures []*signature, options *VerifyOptions) ([]*Verification, error) {
        pipeWriters := make([]*io.PipeWriter, len(signatures))
        // We can't pass pipeWriter to io.MultiWriter directly,
        // we need a slice of io.Writer, but we also need *io.PipeWriter
        // to call Close on it.
        writers := make([]io.Writer, len(signatures))
        chans := make([]chan *Verification, len(signatures))

        for i, sig := range signatures {
                // Be careful with loop variables and goroutines.
                i, sig := i, sig

                chans[i] = make(chan *Verification, 1)

                pr, pw := io.Pipe()
                writers[i] = pw
                pipeWriters[i] = pw

                go func() {
                        v, err := verify(h, pr, h[sig.i], sig.v, options)

                        // Make sure we consume the whole reader, otherwise io.Copy on
                        // other side can block forever.
                        io.Copy(ioutil.Discard, pr)

                        v.Err = err
                        chans[i] <- v
                }()
        }

        if _, err := io.Copy(io.MultiWriter(writers...), r); err != nil {
                return nil, err
        }
        for _, wr := range pipeWriters {
                wr.Close()
        }

        verifications := make([]*Verification, len(signatures))
        for i, ch := range chans {
                verifications[i] = <-ch
        }

        // Return unexpected failures as a separate error.
        for _, v := range verifications {
                err := v.Err
                if err != nil && !IsTempFail(err) && !IsPermFail(err) && !isFail(err) {
                        v.Err = nil
                        return verifications, err
                }
        }
        return verifications, nil
}

func verify(h header, r io.Reader, sigField, sigValue string, options *VerifyOptions) (*Verification, error) {
        verif := new(Verification)

        params, err := parseHeaderParams(sigValue)
        if err != nil {
                return verif, permFailError("malformed signature tags: " + err.Error())
        }

        if params["v"] != "1" {
                return verif, permFailError("incompatible signature version")
        }

        verif.Domain = stripWhitespace(params["d"])

        for _, tag := range requiredTags {
                if _, ok := params[tag]; !ok {
                        return verif, permFailError("signature missing required tag")
                }
        }

        if i, ok := params["i"]; ok {
                verif.Identifier = stripWhitespace(i)
                if !strings.HasSuffix(verif.Identifier, "@"+verif.Domain) && !strings.HasSuffix(verif.Identifier, "."+verif.Domain) {
                        return verif, permFailError("domain mismatch")
                }
        } else {
                verif.Identifier = "@" + verif.Domain
        }

        headerKeys := parseTagList(params["h"])
        ok := false
        for _, k := range headerKeys {
                if strings.EqualFold(k, "from") {
                        ok = true
                        break
                }
        }
        if !ok {
                return verif, permFailError("From field not signed")
        }
        verif.HeaderKeys = headerKeys

        if timeStr, ok := params["t"]; ok {
                t, err := parseTime(timeStr)
                if err != nil {
                        return verif, permFailError("malformed time: " + err.Error())
                }
                verif.Time = t
        }
        if expiresStr, ok := params["x"]; ok {
                t, err := parseTime(expiresStr)
                if err != nil {
                        return verif, permFailError("malformed expiration time: " + err.Error())
                }
                verif.Expiration = t
                if now().After(t) {
                        return verif, permFailError("signature has expired")
                }
        }

        // Query public key
        // TODO: compute hash in parallel
        methods := []string{string(QueryMethodDNSTXT)}
        if methodsStr, ok := params["q"]; ok {
                methods = parseTagList(methodsStr)
        }
        var res *queryResult
        for _, method := range methods {
                if query, ok := queryMethods[QueryMethod(method)]; ok {
                        if options != nil {
                                res, err = query(verif.Domain, stripWhitespace(params["s"]), options.LookupTXT)
                        } else {
                                res, err = query(verif.Domain, stripWhitespace(params["s"]), nil)
                        }
                        break
                }
        }
        if err != nil {
                return verif, err
        } else if res == nil {
                return verif, permFailError("unsupported public key query method")
        }

        // Parse algos
        keyAlgo, hashAlgo, ok := strings.Cut(stripWhitespace(params["a"]), "-")
        if !ok {
                return verif, permFailError("malformed algorithm name")
        }

        // Check hash algo
        if res.HashAlgos != nil {
                ok := false
                for _, algo := range res.HashAlgos {
                        if algo == hashAlgo {
                                ok = true
                                break
                        }
                }
                if !ok {
                        return verif, permFailError("inappropriate hash algorithm")
                }
        }
        var hash crypto.Hash
        switch hashAlgo {
        case "sha1":
                // RFC 8301 section 3.1: rsa-sha1 MUST NOT be used for signing or
                // verifying.
                return verif, permFailError(fmt.Sprintf("hash algorithm too weak: %v", hashAlgo))
        case "sha256":
                hash = crypto.SHA256
        default:
                return verif, permFailError("unsupported hash algorithm")
        }

        // Check key algo
        if res.KeyAlgo != keyAlgo {
                return verif, permFailError("inappropriate key algorithm")
        }

        if res.Services != nil {
                ok := false
                for _, s := range res.Services {
                        if s == "email" {
                                ok = true
                                break
                        }
                }
                if !ok {
                        return verif, permFailError("inappropriate service")
                }
        }

        headerCan, bodyCan := parseCanonicalization(params["c"])
        if _, ok := canonicalizers[headerCan]; !ok {
                return verif, permFailError("unsupported header canonicalization algorithm")
        }
        if _, ok := canonicalizers[bodyCan]; !ok {
                return verif, permFailError("unsupported body canonicalization algorithm")
        }

        // The body length "l" parameter is insecure, because it allows parts of
        // the message body to not be signed. Reject messages which have it set.
        if _, ok := params["l"]; ok {
                // TODO: technically should be policyError
                return verif, failError("message contains an insecure body length tag")
        }

        // Parse body hash and signature
        bodyHashed, err := decodeBase64String(params["bh"])
        if err != nil {
                return verif, permFailError("malformed body hash: " + err.Error())
        }
        sig, err := decodeBase64String(params["b"])
        if err != nil {
                return verif, permFailError("malformed signature: " + err.Error())
        }

        // Check body hash
        hasher := hash.New()
        wc := canonicalizers[bodyCan].CanonicalizeBody(hasher)
        if _, err := io.Copy(wc, r); err != nil {
                return verif, err
        }
        if err := wc.Close(); err != nil {
                return verif, err
        }
        if subtle.ConstantTimeCompare(hasher.Sum(nil), bodyHashed) != 1 {
                return verif, failError("body hash did not verify")
        }

        // Compute data hash
        hasher.Reset()
        picker := newHeaderPicker(h)
        for _, key := range headerKeys {
                kv := picker.Pick(key)
                if kv == "" {
                        // The field MAY contain names of header fields that do not exist
                        // when signed; nonexistent header fields do not contribute to the
                        // signature computation
                        continue
                }

                kv = canonicalizers[headerCan].CanonicalizeHeader(kv)
                if _, err := hasher.Write([]byte(kv)); err != nil {
                        return verif, err
                }
        }
        canSigField := removeSignature(sigField)
        canSigField = canonicalizers[headerCan].CanonicalizeHeader(canSigField)
        canSigField = strings.TrimRight(canSigField, "\r\n")
        if _, err := hasher.Write([]byte(canSigField)); err != nil {
                return verif, err
        }
        hashed := hasher.Sum(nil)

        // Check signature
        if err := res.Verifier.Verify(hash, hashed, sig); err != nil {
                return verif, failError("signature did not verify: " + err.Error())
        }

        return verif, nil
}

func parseTagList(s string) []string {
        tags := strings.Split(s, ":")
        for i, t := range tags {
                tags[i] = stripWhitespace(t)
        }
        return tags
}

func parseCanonicalization(s string) (headerCan, bodyCan Canonicalization) {
        headerCan = CanonicalizationSimple
        bodyCan = CanonicalizationSimple

        cans := strings.SplitN(stripWhitespace(s), "/", 2)
        if cans[0] != "" {
                headerCan = Canonicalization(cans[0])
        }
        if len(cans) > 1 {
                bodyCan = Canonicalization(cans[1])
        }
        return
}

func parseTime(s string) (time.Time, error) {
        sec, err := strconv.ParseInt(stripWhitespace(s), 10, 64)
        if err != nil {
                return time.Time{}, err
        }
        return time.Unix(sec, 0), nil
}

func decodeBase64String(s string) ([]byte, error) {
        return base64.StdEncoding.DecodeString(stripWhitespace(s))
}

func stripWhitespace(s string) string {
        return strings.Map(func(r rune) rune {
                if unicode.IsSpace(r) {
                        return -1
                }
                return r
        }, s)
}

var sigRegex = regexp.MustCompile(`(b\s*=)[^;]+`)

func removeSignature(s string) string {
        return sigRegex.ReplaceAllString(s, "$1")
}

package dmarc

import (
        "errors"
        "fmt"
        "net"
        "strconv"
        "strings"
        "time"
)

type tempFailError string

func (err tempFailError) Error() string {
        return "dmarc: " + string(err)
}

// IsTempFail returns true if the error returned by Lookup is a temporary
// failure.
func IsTempFail(err error) bool {
        _, ok := err.(tempFailError)
        return ok
}

var ErrNoPolicy = errors.New("dmarc: no policy found for domain")

// LookupOptions allows to customize the default signature verification behavior
// LookupTXT returns the DNS TXT records for the given domain name. If nil, net.LookupTXT is used
type LookupOptions struct {
        LookupTXT func(domain string) ([]string, error)
}

// Lookup queries a DMARC record for a specified domain.
func Lookup(domain string) (*Record, error) {
        return LookupWithOptions(domain, nil)
}

func LookupWithOptions(domain string, options *LookupOptions) (*Record, error) {
        var txts []string
        var err error
        if options != nil && options.LookupTXT != nil {
                txts, err = options.LookupTXT("_dmarc." + domain)
        } else {
                txts, err = net.LookupTXT("_dmarc." + domain)
        }
        if netErr, ok := err.(net.Error); ok && netErr.Temporary() {
                return nil, tempFailError("TXT record unavailable: " + err.Error())
        } else if err != nil {
                if dnsErr, ok := err.(*net.DNSError); ok && dnsErr.IsNotFound {
                        return nil, ErrNoPolicy
                }
                return nil, errors.New("dmarc: failed to lookup TXT record: " + err.Error())
        }
        if len(txts) == 0 {
                return nil, ErrNoPolicy
        }

        // Long keys are split in multiple parts
        txt := strings.Join(txts, "")
        return Parse(txt)
}

func Parse(txt string) (*Record, error) {
        params, err := parseParams(txt)
        if err != nil {
                return nil, err
        }

        if params["v"] != "DMARC1" {
                return nil, errors.New("dmarc: unsupported DMARC version")
        }

        rec := new(Record)

        p, ok := params["p"]
        if !ok {
                return nil, errors.New("dmarc: record is missing a 'p' parameter")
        }
        rec.Policy, err = parsePolicy(p, "p")
        if err != nil {
                return nil, err
        }

        rec.DKIMAlignment = AlignmentRelaxed
        if adkim, ok := params["adkim"]; ok {
                rec.DKIMAlignment, err = parseAlignmentMode(adkim, "adkim")
                if err != nil {
                        return nil, err
                }
        }

        rec.SPFAlignment = AlignmentRelaxed
        if aspf, ok := params["aspf"]; ok {
                rec.SPFAlignment, err = parseAlignmentMode(aspf, "aspf")
                if err != nil {
                        return nil, err
                }
        }

        if fo, ok := params["fo"]; ok {
                rec.FailureOptions, err = parseFailureOptions(fo)
                if err != nil {
                        return nil, err
                }
        }

        if pct, ok := params["pct"]; ok {
                i, err := strconv.Atoi(pct)
                if err != nil {
                        return nil, fmt.Errorf("dmarc: invalid parameter 'pct': %v", err)
                }
                if i < 0 || i > 100 {
                        return nil, fmt.Errorf("dmarc: invalid parameter 'pct': value %v out of bounds", i)
                }
                rec.Percent = &i
        }

        if rf, ok := params["rf"]; ok {
                l := strings.Split(rf, ":")
                rec.ReportFormat = make([]ReportFormat, len(l))
                for i, f := range l {
                        switch f {
                        case "afrf":
                                rec.ReportFormat[i] = ReportFormat(f)
                        default:
                                return nil, errors.New("dmarc: invalid parameter 'rf'")
                        }
                }
        }

        if ri, ok := params["ri"]; ok {
                i, err := strconv.Atoi(ri)
                if err != nil {
                        return nil, fmt.Errorf("dmarc: invalid parameter 'ri': %v", err)
                }
                if i <= 0 {
                        return nil, fmt.Errorf("dmarc: invalid parameter 'ri': negative or zero duration")
                }
                rec.ReportInterval = time.Duration(i) * time.Second
        }

        if rua, ok := params["rua"]; ok {
                rec.ReportURIAggregate = parseURIList(rua)
        }

        if ruf, ok := params["ruf"]; ok {
                rec.ReportURIFailure = parseURIList(ruf)
        }

        if sp, ok := params["sp"]; ok {
                rec.SubdomainPolicy, err = parsePolicy(sp, "sp")
                if err != nil {
                        return nil, err
                }
        }

        return rec, nil
}

func parseParams(s string) (map[string]string, error) {
        pairs := strings.Split(s, ";")
        params := make(map[string]string)
        for _, s := range pairs {
                kv := strings.SplitN(s, "=", 2)
                if len(kv) != 2 {
                        if strings.TrimSpace(s) == "" {
                                continue
                        }
                        return params, errors.New("dmarc: malformed params")
                }

                params[strings.TrimSpace(kv[0])] = strings.TrimSpace(kv[1])
        }
        return params, nil
}

func parsePolicy(s, param string) (Policy, error) {
        switch s {
        case "none", "quarantine", "reject":
                return Policy(s), nil
        default:
                return "", fmt.Errorf("dmarc: invalid policy for parameter '%v'", param)
        }
}

func parseAlignmentMode(s, param string) (AlignmentMode, error) {
        switch s {
        case "r", "s":
                return AlignmentMode(s), nil
        default:
                return "", fmt.Errorf("dmarc: invalid alignment mode for parameter '%v'", param)
        }
}

func parseFailureOptions(s string) (FailureOptions, error) {
        l := strings.Split(s, ":")
        var opts FailureOptions
        for _, o := range l {
                switch strings.TrimSpace(o) {
                case "0":
                        opts |= FailureAll
                case "1":
                        opts |= FailureAny
                case "d":
                        opts |= FailureDKIM
                case "s":
                        opts |= FailureSPF
                default:
                        return 0, errors.New("dmarc: invalid failure option in parameter 'fo'")
                }
        }
        return opts, nil
}

func parseURIList(s string) []string {
        l := strings.Split(s, ",")
        for i, u := range l {
                l[i] = strings.TrimSpace(u)
        }
        return l
}