From 9e814d001ad2588f13048beb14f7a7dc07155ebc Mon Sep 17 00:00:00 2001
From: Jonathan Giddy <jgiddy@mesosphere.com>
Date: Thu, 11 Apr 2019 06:58:38 +0100
Subject: [PATCH] Escape special characters in Prometheus output

Co-authored-by: Furhan Shabir <iamfurhan@gmail.com>
---
 ..._vhost_traffic_status_display_prometheus.c | 26 +++--
 src/ngx_http_vhost_traffic_status_string.c    | 96 +++++++++++++++++++
 src/ngx_http_vhost_traffic_status_string.h    |  3 +-
 3 files changed, 114 insertions(+), 11 deletions(-)

diff --git a/src/ngx_http_vhost_traffic_status_display_prometheus.c b/src/ngx_http_vhost_traffic_status_display_prometheus.c
index fac578e..9696e0b 100644
--- a/src/ngx_http_vhost_traffic_status_display_prometheus.c
+++ b/src/ngx_http_vhost_traffic_status_display_prometheus.c
@@ -125,7 +125,7 @@ u_char *
 ngx_http_vhost_traffic_status_display_prometheus_set_server(ngx_http_request_t *r,
     u_char *buf, ngx_rbtree_node_t *node)
 {
-    ngx_str_t                                  key;
+    ngx_str_t                                  key, escaped_key;
     ngx_http_vhost_traffic_status_ctx_t       *ctx;
     ngx_http_vhost_traffic_status_node_t      *vtsn;
     ngx_http_vhost_traffic_status_loc_conf_t  *vtscf;
@@ -141,7 +141,8 @@ ngx_http_vhost_traffic_status_display_prometheus_set_server(ngx_http_request_t *
             key.data = vtsn->data;
             key.len = vtsn->len;
 
-            buf = ngx_http_vhost_traffic_status_display_prometheus_set_server_node(r, buf, &key, vtsn);
+            ngx_http_vhost_traffic_status_escape_prometheus(r->pool, &escaped_key, key.data, key.len);
+            buf = ngx_http_vhost_traffic_status_display_prometheus_set_server_node(r, buf, &escaped_key, vtsn);
 
             /* calculates the sum */
             vtscf->stats.stat_request_counter += vtsn->stat_request_counter;
@@ -268,7 +269,7 @@ u_char *
 ngx_http_vhost_traffic_status_display_prometheus_set_filter(ngx_http_request_t *r,
     u_char *buf, ngx_rbtree_node_t *node)
 {
-    ngx_str_t                              key;
+    ngx_str_t                              key, escaped_key;
     ngx_http_vhost_traffic_status_ctx_t   *ctx;
     ngx_http_vhost_traffic_status_node_t  *vtsn;
 
@@ -281,7 +282,8 @@ ngx_http_vhost_traffic_status_display_prometheus_set_filter(ngx_http_request_t *
             key.data = vtsn->data;
             key.len = vtsn->len;
 
-            buf = ngx_http_vhost_traffic_status_display_prometheus_set_filter_node(r, buf, &key, vtsn);
+            ngx_http_vhost_traffic_status_escape_prometheus(r->pool, &escaped_key, key.data, key.len);
+            buf = ngx_http_vhost_traffic_status_display_prometheus_set_filter_node(r, buf, &escaped_key, vtsn);
         }
 
         buf = ngx_http_vhost_traffic_status_display_prometheus_set_filter(r, buf, node->left);
@@ -388,7 +390,7 @@ u_char *
 ngx_http_vhost_traffic_status_display_prometheus_set_upstream(ngx_http_request_t *r,
     u_char *buf, ngx_rbtree_node_t *node)
 {
-    ngx_str_t                              key;
+    ngx_str_t                              key, escaped_key;
     ngx_http_vhost_traffic_status_ctx_t   *ctx;
     ngx_http_vhost_traffic_status_node_t  *vtsn;
 
@@ -403,7 +405,8 @@ ngx_http_vhost_traffic_status_display_prometheus_set_upstream(ngx_http_request_t
             key.data = vtsn->data;
             key.len = vtsn->len;
 
-            buf = ngx_http_vhost_traffic_status_display_prometheus_set_upstream_node(r, buf, &key, vtsn);
+            ngx_http_vhost_traffic_status_escape_prometheus(r->pool, &escaped_key, key.data, key.len);
+            buf = ngx_http_vhost_traffic_status_display_prometheus_set_upstream_node(r, buf, &escaped_key, vtsn);
         }
 
         buf = ngx_http_vhost_traffic_status_display_prometheus_set_upstream(r, buf, node->left);
@@ -450,7 +453,7 @@ u_char *
 ngx_http_vhost_traffic_status_display_prometheus_set_cache(ngx_http_request_t *r,
     u_char *buf, ngx_rbtree_node_t *node)
 {
-    ngx_str_t                              key;
+    ngx_str_t                              key, escaped_key;
     ngx_http_vhost_traffic_status_ctx_t   *ctx;
     ngx_http_vhost_traffic_status_node_t  *vtsn;
 
@@ -463,7 +466,8 @@ ngx_http_vhost_traffic_status_display_prometheus_set_cache(ngx_http_request_t *r
             key.data = vtsn->data;
             key.len = vtsn->len;
 
-            buf = ngx_http_vhost_traffic_status_display_prometheus_set_cache_node(r, buf, &key, vtsn);
+            ngx_http_vhost_traffic_status_escape_prometheus(r->pool, &escaped_key, key.data, key.len);
+            buf = ngx_http_vhost_traffic_status_display_prometheus_set_cache_node(r, buf, &escaped_key, vtsn);
         }
 
         buf = ngx_http_vhost_traffic_status_display_prometheus_set_cache(r, buf, node->left);
@@ -480,6 +484,7 @@ u_char *
 ngx_http_vhost_traffic_status_display_prometheus_set(ngx_http_request_t *r,
     u_char *buf)
 {
+    ngx_str_t                                 escaped_key;
     u_char                                    *o, *s;
     ngx_rbtree_node_t                         *node;
     ngx_http_vhost_traffic_status_ctx_t       *ctx;
@@ -505,8 +510,9 @@ ngx_http_vhost_traffic_status_display_prometheus_set(ngx_http_request_t *r,
 #endif
     buf = ngx_http_vhost_traffic_status_display_prometheus_set_server(r, buf, node);
 
-    buf = ngx_http_vhost_traffic_status_display_prometheus_set_server_node(r, buf, &vtscf->sum_key,
-                                                                           &vtscf->stats);
+    ngx_http_vhost_traffic_status_escape_prometheus(r->pool, &escaped_key, vtscf->sum_key.data, vtscf->sum_key.len);
+    buf = ngx_http_vhost_traffic_status_display_prometheus_set_server_node(r, buf, &escaped_key, &vtscf->stats);
+    
     /* filterZones */
     o = buf;
 
diff --git a/src/ngx_http_vhost_traffic_status_string.c b/src/ngx_http_vhost_traffic_status_string.c
index 800c058..d2195f2 100644
--- a/src/ngx_http_vhost_traffic_status_string.c
+++ b/src/ngx_http_vhost_traffic_status_string.c
@@ -3,6 +3,7 @@
  * Copyright (C) YoungJoo Kim (vozlt)
  */
 
+#include <ctype.h>
 
 #include <ngx_config.h>
 #include <ngx_core.h>
@@ -164,4 +165,99 @@ ngx_http_vhost_traffic_status_replace_strc(ngx_str_t *buf,
     return NGX_OK;
 }
 
+
+ngx_int_t
+ngx_http_vhost_traffic_status_escape_prometheus(ngx_pool_t *pool, ngx_str_t *buf, u_char *p, size_t n)
+{
+    u_char  c, *pa, *pb, *last, *char_end;
+    size_t  size;
+    u_char  HEX_MAP[] = "0123456789ABCDEF";
+
+    last = p + n;
+    pa = p;
+    size = 0;
+
+    /* Find the first character that needs to be escaped */
+    while (pa < last) {
+        if isascii(*pa) {
+            if (*pa == '"' || *pa == '\\' || *pa == '\n') {
+                break;
+            } else {
+                pa++;
+            }
+        } else {
+            char_end = pa;
+            if (ngx_utf8_decode(&char_end, last - pa) > 0x10ffff) {
+                break;
+            } else {
+                pa = char_end;
+            }
+        }
+    }
+
+    if (pa == last) {
+        // no escapes required - return the original string
+        buf->data = p;
+        buf->len = n;
+        return NGX_OK;
+    }
+
+    size = pa - p;
+
+    /* Allocate enough space for the unescaped prefix and worst case for remainder */
+    buf->data = ngx_pcalloc(pool, size + (n - size) * 5);
+    if (buf->data == NULL) {
+        /*
+            Return the unescaped string up to the first special character 
+            in case the caller does not handle the error.
+        */
+        buf->data = p;
+        buf->len = size;
+        return NGX_ERROR;
+    }
+
+    /* Copy `size` unescaped characters to start of destination. */
+    pb = ngx_copy(buf->data, p, size);
+
+    /* Individually copy remaining characters to destination, escaping as necessary */
+    while (pa < last) {
+        if (isascii(*pa)) {
+            if (*pa == '"' || *pa == '\\') {
+                *pb++ = '\\';
+                *pb++ = *pa++;
+                size += 2;
+            } else if (*pa == '\n') {
+                *pb++ = '\\';
+                *pb++ = 'n';
+                pa++;
+                size += 2;
+            } else {
+                *pb++ = *pa++;
+                size++;
+            }
+        } else {
+            char_end = pa;
+            if (ngx_utf8_decode(&char_end, last - pa) > 0x10ffff) {
+                /* invalid UTF-8 - escape single char to allow resynchronization */
+                c = *pa++;
+                /* two slashes are required to be valid encoding for prometheus*/
+                *pb++ = '\\';
+                *pb++ = '\\';
+                *pb++ = 'x';
+                *pb++ = HEX_MAP[c >> 4];
+                *pb++ = HEX_MAP[c & 0x0f];
+                size += 5;
+            } else {
+                while (pa < char_end) {
+                    *pb++ = *pa++;
+                    size++;
+                }
+            }
+        }
+    }
+
+    buf->len = size;
+    return NGX_OK;
+}
+
 /* vi:set ft=c ts=4 sw=4 et fdm=marker: */
diff --git a/src/ngx_http_vhost_traffic_status_string.h b/src/ngx_http_vhost_traffic_status_string.h
index 20f3f4b..d8f4b36 100644
--- a/src/ngx_http_vhost_traffic_status_string.h
+++ b/src/ngx_http_vhost_traffic_status_string.h
@@ -19,7 +19,8 @@ ngx_int_t ngx_http_vhost_traffic_status_replace_chrc(ngx_str_t *buf,
     u_char in, u_char to);
 ngx_int_t ngx_http_vhost_traffic_status_replace_strc(ngx_str_t *buf,
     ngx_str_t *dst, u_char c);
-
+ngx_int_t ngx_http_vhost_traffic_status_escape_prometheus(ngx_pool_t *pool, ngx_str_t *buf, 
+	u_char *p, size_t n);
 
 #endif /* _NGX_HTTP_VTS_STRING_H_INCLUDED_ */
 
