1 files changed, 106 insertions, 0 deletions
diff --git a/att script/4(v6 DDoS)/code/辅助权威服务器/plugin/prometheus/setup.go b/att script/4(v6 DDoS)/code/辅助权威服务器/plugin/prometheus/setup.go
new file mode 100644
index 0000000..e5ca825
--- /dev/null
+++ b/att script/4(v6 DDoS)/code/辅助权威服务器/plugin/prometheus/setup.go
@@ -0,0 +1,106 @@
+package prometheus
+
+import (
+	"net"
+	"ohmydns2/core/dnsserver"
+	"ohmydns2/ohmain"
+	"ohmydns2/plugin"
+	olog "ohmydns2/plugin/pkg/log"
+	"ohmydns2/plugin/pkg/uniq"
+	"ohmydns2/plugin/prometheus/vars"
+
+	"github.com/coredns/caddy"
+
+	"runtime"
+)
+
+var (
+	log      = olog.NewWithPlugin("prometheus")
+	u        = uniq.New()
+	registry = newReg()
+)
+
+func init() { plugin.Register("prometheus", setup) }
+
+func setup(c *caddy.Controller) error {
+	m, err := parse(c)
+	if err != nil {
+		return plugin.Error("prometheus", err)
+	}
+	m.Reg = registry.getOrSet(m.Addr, m.Reg)
+
+	c.OnStartup(func() error { m.Reg = registry.getOrSet(m.Addr, m.Reg); u.Set(m.Addr, m.OnStartup); return nil })
+	c.OnRestartFailed(func() error { m.Reg = registry.getOrSet(m.Addr, m.Reg); u.Set(m.Addr, m.OnStartup); return nil })
+
+	c.OnStartup(func() error { return u.ForEach() })
+	c.OnRestartFailed(func() error { return u.ForEach() })
+
+	c.OnStartup(func() error {
+		conf := dnsserver.GetConfig(c)
+		for _, h := range conf.ListenHosts {
+			addrstr := conf.Transport + "://" + net.JoinHostPort(h, conf.Port)
+			for _, p := range conf.Handlers() {
+				vars.PluginEnabled.WithLabelValues(addrstr, conf.Zone, conf.ViewName, p.Name()).Set(1)
+			}
+		}
+		return nil
+	})
+	c.OnRestartFailed(func() error {
+		conf := dnsserver.GetConfig(c)
+		for _, h := range conf.ListenHosts {
+			addrstr := conf.Transport + "://" + net.JoinHostPort(h, conf.Port)
+			for _, p := range conf.Handlers() {
+				vars.PluginEnabled.WithLabelValues(addrstr, conf.Zone, conf.ViewName, p.Name()).Set(1)
+			}
+		}
+		return nil
+	})
+
+	c.OnRestart(m.OnRestart)
+	c.OnRestart(func() error { vars.PluginEnabled.Reset(); return nil })
+	c.OnFinalShutdown(m.OnFinalShutdown)
+
+	// Initialize metrics.
+	buildInfo.WithLabelValues(ohmain.OMVersion, ohmain.GitCommit, runtime.Version()).Set(1)
+
+	dnsserver.GetConfig(c).AddPlugin(func(next plugin.Handler) plugin.Handler {
+		m.Next = next
+		return m
+	})
+
+	return nil
+}
+
+func parse(c *caddy.Controller) (*Metrics, error) {
+	met := New(defaultAddr)
+
+	i := 0
+	for c.Next() {
+		if i > 0 {
+			return nil, plugin.ErrOnce
+		}
+		i++
+
+		zones := plugin.OriginsFromArgsOrServerBlock(nil /* args */, c.ServerBlockKeys)
+		for _, z := range zones {
+			met.AddZone(z)
+		}
+		args := c.RemainingArgs()
+
+		switch len(args) {
+		case 0:
+		case 1:
+			met.Addr = args[0]
+			_, _, e := net.SplitHostPort(met.Addr)
+			if e != nil {
+				return met, e
+			}
+		default:
+			return met, c.ArgErr()
+		}
+	}
+	return met, nil
+}
+
+// defaultAddr is the address the where the metrics are exported by default.
+const defaultAddr = "localhost:9153"