-
Notifications
You must be signed in to change notification settings - Fork 2.1k
/
cancel.go
66 lines (52 loc) · 1.47 KB
/
cancel.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
// Package cancel implements a plugin adds a canceling context to each request.
package cancel
import (
"context"
"fmt"
"time"
"github.com/coredns/caddy"
"github.com/coredns/coredns/core/dnsserver"
"github.com/coredns/coredns/plugin"
"github.com/miekg/dns"
)
func init() { plugin.Register("cancel", setup) }
func setup(c *caddy.Controller) error {
ca := Cancel{}
for c.Next() {
args := c.RemainingArgs()
switch len(args) {
case 0:
ca.timeout = 5001 * time.Millisecond
case 1:
dur, err := time.ParseDuration(args[0])
if err != nil {
return plugin.Error("cancel", fmt.Errorf("invalid duration: %q", args[0]))
}
if dur <= 0 {
return plugin.Error("cancel", fmt.Errorf("invalid negative duration: %q", args[0]))
}
ca.timeout = dur
default:
return plugin.Error("cancel", c.ArgErr())
}
}
dnsserver.GetConfig(c).AddPlugin(func(next plugin.Handler) plugin.Handler {
ca.Next = next
return ca
})
return nil
}
// Cancel is a plugin that adds a canceling context to each request's context.
type Cancel struct {
timeout time.Duration
Next plugin.Handler
}
// ServeDNS implements the plugin.Handler interface.
func (c Cancel) ServeDNS(ctx context.Context, w dns.ResponseWriter, r *dns.Msg) (int, error) {
ctx, cancel := context.WithTimeout(ctx, c.timeout)
code, err := plugin.NextOrFailure(c.Name(), c.Next, ctx, w, r)
cancel()
return code, err
}
// Name implements the Handler interface.
func (c Cancel) Name() string { return "cancel" }