From ded7381cb0b644677c96ec28dcd73609433eb21e Mon Sep 17 00:00:00 2001
From: Daniel Czerwonk <daniel@dan-nrw.de>
Date: Sun, 11 Dec 2016 16:26:57 +0100
Subject: [PATCH] added better error handling, ipv6 and ipv4 can now be
 disabled on demand

---
 bird_client.go | 47 ++++++++++++++++++++++++++++-------------------
 main.go        | 43 ++++++++++++++++++++++++++++++++++++-------
 parser.go      |  2 +-
 3 files changed, 65 insertions(+), 27 deletions(-)

diff --git a/bird_client.go b/bird_client.go
index fdab9e3..308d4a0 100644
--- a/bird_client.go
+++ b/bird_client.go
@@ -17,36 +17,45 @@ along with this program.  If not, see <http://www.gnu.org/licenses/>
 
 package main
 
-import (
-	"log"
-	"os/exec"
-)
+import "os/exec"
 
-func getSessions() []*session {
-	birdSessions := getSessionsFromBird(4)
-	bird6Sessions := getSessionsFromBird(6)
+func getSessions() ([]*session, error) {
+	sessions := make([]*session, 0)
 
-	return append(birdSessions, bird6Sessions...)
+	if *birdEnabled {
+		s, err := getSessionsFromBird(4)
+		if err != nil {
+			return nil, err
+		}
+		sessions = append(sessions, s...)
+	}
+
+	if *bird6Enabled {
+		s, err := getSessionsFromBird(6)
+		if err != nil {
+			return nil, err
+		}
+		sessions = append(sessions, s...)
+	}
+
+	return sessions, nil
 }
 
-func getSessionsFromBird(ipVersion int) []*session {
+func getSessionsFromBird(ipVersion int) ([]*session, error) {
 	client := *birdClient
 
 	if ipVersion == 6 {
 		client += "6"
 	}
 
-	output := getBirdOutput(client)
-	return parseOutput(output, ipVersion)
-}
-
-func getBirdOutput(birdClient string) []byte {
-	b, err := exec.Command(birdClient, "show", "protocols", "all").Output()
-
+	output, err := getBirdOutput(client)
 	if err != nil {
-		b = make([]byte, 0)
-		log.Println(err)
+		return nil, err
 	}
 
-	return b
+	return parseOutput(output, ipVersion), nil
+}
+
+func getBirdOutput(birdClient string) ([]byte, error) {
+	return exec.Command(birdClient, "show", "protocols", "all").Output()
 }
diff --git a/main.go b/main.go
index 67da5f5..d6126d0 100644
--- a/main.go
+++ b/main.go
@@ -18,6 +18,8 @@ along with this program.  If not, see <http://www.gnu.org/licenses/>
 package main
 
 import (
+	"bufio"
+	"bytes"
 	"flag"
 	"fmt"
 	"io"
@@ -35,13 +37,15 @@ type session struct {
 	uptime      int
 }
 
-const version string = "0.2.1"
+const version string = "0.3"
 
 var (
 	showVersion   = flag.Bool("version", false, "Print version information.")
 	listenAddress = flag.String("web.listen-address", ":9200", "Address on which to expose metrics and web interface.")
 	metricsPath   = flag.String("web.telemetry-path", "/metrics", "Path under which to expose metrics.")
 	birdClient    = flag.String("bird.client", "birdc", "Binary to communicate with the bird routing daemon")
+	birdEnabled   = flag.Bool("bird.ipv4", true, "Get sessions from bird")
+	bird6Enabled  = flag.Bool("bird.ipv6", true, "Get sessions from bird6")
 )
 
 func main() {
@@ -64,21 +68,46 @@ func printVersion() {
 
 func startServer() {
 	fmt.Printf("Starting bird BGP exporter (Version: %s)\n", version)
-	http.HandleFunc(*metricsPath, handleMetricsRequest)
+	http.HandleFunc(*metricsPath, errorHandler(handleMetricsRequest))
 
 	fmt.Printf("Listening for %s on %s\n", *metricsPath, *listenAddress)
 	log.Fatal(http.ListenAndServe(*listenAddress, nil))
 }
 
-func handleMetricsRequest(w http.ResponseWriter, r *http.Request) {
-	sessions := getSessions()
+func errorHandler(f func(io.Writer, *http.Request) error) http.HandlerFunc {
+	return func(w http.ResponseWriter, r *http.Request) {
+		var buf bytes.Buffer
+		wr := bufio.NewWriter(&buf)
+		err := f(wr, r)
+		wr.Flush()
 
-	for _, s := range sessions {
-		writeLineForSession(s, w)
+		if err != nil {
+			log.Println(err)
+			http.Error(w, err.Error(), http.StatusInternalServerError)
+		}
+
+		_, err = w.Write(buf.Bytes())
+
+		if err != nil {
+			log.Println(err)
+		}
 	}
 }
 
-func writeLineForSession(s *session, w io.Writer) {
+func handleMetricsRequest(w io.Writer, r *http.Request) error {
+	sessions, err := getSessions()
+	if err != nil {
+		return err
+	}
+
+	for _, s := range sessions {
+		writeForSession(s, w)
+	}
+
+	return nil
+}
+
+func writeForSession(s *session, w io.Writer) {
 	fmt.Fprintf(w, "bgp%d_session_up{name=\"%s\"} %d\n", s.ipVersion, s.name, s.established)
 	fmt.Fprintf(w, "bgp%d_session_prefix_count_import{name=\"%s\"} %d\n", s.ipVersion, s.name, s.imported)
 	fmt.Fprintf(w, "bgp%d_session_prefix_count_export{name=\"%s\"} %d\n", s.ipVersion, s.name, s.exported)
diff --git a/parser.go b/parser.go
index d7f346e..e2ae979 100644
--- a/parser.go
+++ b/parser.go
@@ -48,7 +48,7 @@ func parseOutput(data []byte, ipVersion int) []*session {
 
 	for scanner.Scan() {
 		line := scanner.Text()
-		if session, res := parseLineForSession(line, ipVersion); res == true {
+		if session, ok := parseLineForSession(line, ipVersion); ok {
 			current = session
 			sessions = append(sessions, current)
 		}