canonical.go raw
1 package dkim
2
3 import (
4 "io"
5 "strings"
6 )
7
8 // Canonicalization is a canonicalization algorithm.
9 type Canonicalization string
10
11 const (
12 CanonicalizationSimple Canonicalization = "simple"
13 CanonicalizationRelaxed = "relaxed"
14 )
15
16 type canonicalizer interface {
17 CanonicalizeHeader(s string) string
18 CanonicalizeBody(w io.Writer) io.WriteCloser
19 }
20
21 var canonicalizers = map[Canonicalization]canonicalizer{
22 CanonicalizationSimple: new(simpleCanonicalizer),
23 CanonicalizationRelaxed: new(relaxedCanonicalizer),
24 }
25
26 // crlfFixer fixes any lone LF without a preceding CR.
27 type crlfFixer struct {
28 cr bool
29 }
30
31 func (cf *crlfFixer) Fix(b []byte) []byte {
32 res := make([]byte, 0, len(b))
33 for _, ch := range b {
34 prevCR := cf.cr
35 cf.cr = false
36 switch ch {
37 case '\r':
38 cf.cr = true
39 case '\n':
40 if !prevCR {
41 res = append(res, '\r')
42 }
43 }
44 res = append(res, ch)
45 }
46 return res
47 }
48
49 type simpleCanonicalizer struct{}
50
51 func (c *simpleCanonicalizer) CanonicalizeHeader(s string) string {
52 return s
53 }
54
55 type simpleBodyCanonicalizer struct {
56 w io.Writer
57 crlfBuf []byte
58 crlfFixer crlfFixer
59 }
60
61 func (c *simpleBodyCanonicalizer) Write(b []byte) (int, error) {
62 written := len(b)
63 b = append(c.crlfBuf, b...)
64
65 b = c.crlfFixer.Fix(b)
66
67 end := len(b)
68 // If it ends with \r, maybe the next write will begin with \n
69 if end > 0 && b[end-1] == '\r' {
70 end--
71 }
72 // Keep all \r\n sequences
73 for end >= 2 {
74 prev := b[end-2]
75 cur := b[end-1]
76 if prev != '\r' || cur != '\n' {
77 break
78 }
79 end -= 2
80 }
81
82 c.crlfBuf = b[end:]
83
84 var err error
85 if end > 0 {
86 _, err = c.w.Write(b[:end])
87 }
88 return written, err
89 }
90
91 func (c *simpleBodyCanonicalizer) Close() error {
92 // Flush crlfBuf if it ends with a single \r (without a matching \n)
93 if len(c.crlfBuf) > 0 && c.crlfBuf[len(c.crlfBuf)-1] == '\r' {
94 if _, err := c.w.Write(c.crlfBuf); err != nil {
95 return err
96 }
97 }
98 c.crlfBuf = nil
99
100 if _, err := c.w.Write([]byte(crlf)); err != nil {
101 return err
102 }
103 return nil
104 }
105
106 func (c *simpleCanonicalizer) CanonicalizeBody(w io.Writer) io.WriteCloser {
107 return &simpleBodyCanonicalizer{w: w}
108 }
109
110 type relaxedCanonicalizer struct{}
111
112 func (c *relaxedCanonicalizer) CanonicalizeHeader(s string) string {
113 k, v, ok := strings.Cut(s, ":")
114 if !ok {
115 return strings.TrimSpace(strings.ToLower(s)) + ":" + crlf
116 }
117
118 k = strings.TrimSpace(strings.ToLower(k))
119 v = strings.Join(strings.FieldsFunc(v, func(r rune) bool {
120 return r == ' ' || r == '\t' || r == '\n' || r == '\r'
121 }), " ")
122 return k + ":" + v + crlf
123 }
124
125 type relaxedBodyCanonicalizer struct {
126 w io.Writer
127 crlfBuf []byte
128 wsp bool
129 written bool
130 crlfFixer crlfFixer
131 }
132
133 func (c *relaxedBodyCanonicalizer) Write(b []byte) (int, error) {
134 written := len(b)
135
136 b = c.crlfFixer.Fix(b)
137
138 canonical := make([]byte, 0, len(b))
139 for _, ch := range b {
140 if ch == ' ' || ch == '\t' {
141 c.wsp = true
142 } else if ch == '\r' || ch == '\n' {
143 c.wsp = false
144 c.crlfBuf = append(c.crlfBuf, ch)
145 } else {
146 if len(c.crlfBuf) > 0 {
147 canonical = append(canonical, c.crlfBuf...)
148 c.crlfBuf = c.crlfBuf[:0]
149 }
150 if c.wsp {
151 canonical = append(canonical, ' ')
152 c.wsp = false
153 }
154
155 canonical = append(canonical, ch)
156 }
157 }
158
159 if !c.written && len(canonical) > 0 {
160 c.written = true
161 }
162
163 _, err := c.w.Write(canonical)
164 return written, err
165 }
166
167 func (c *relaxedBodyCanonicalizer) Close() error {
168 if c.written {
169 if _, err := c.w.Write([]byte(crlf)); err != nil {
170 return err
171 }
172 }
173 return nil
174 }
175
176 func (c *relaxedCanonicalizer) CanonicalizeBody(w io.Writer) io.WriteCloser {
177 return &relaxedBodyCanonicalizer{w: w}
178 }
179
180 type limitedWriter struct {
181 W io.Writer
182 N int64
183 }
184
185 func (w *limitedWriter) Write(b []byte) (int, error) {
186 if w.N <= 0 {
187 return len(b), nil
188 }
189
190 skipped := 0
191 if int64(len(b)) > w.N {
192 b = b[:w.N]
193 skipped = int(int64(len(b)) - w.N)
194 }
195
196 n, err := w.W.Write(b)
197 w.N -= int64(n)
198 return n + skipped, err
199 }
200