1 /** @file
2 
3     Plugin to perform background fetches of certain content that would
4     otherwise not be cached. For example, Range: requests / responses.
5 
6     @section license License
7 
8     Licensed to the Apache Software Foundation (ASF) under one
9     or more contributor license agreements.  See the NOTICE file
10     distributed with this work for additional information
11     regarding copyright ownership.  The ASF licenses this file
12     to you under the Apache License, Version 2.0 (the
13     "License"); you may not use this file except in compliance
14     with the License.  You may obtain a copy of the License at
15 
16     http://www.apache.org/licenses/LICENSE-2.0
17 
18     Unless required by applicable law or agreed to in writing, software
19     distributed under the License is distributed on an "AS IS" BASIS,
20     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
21     See the License for the specific language governing permissions and
22     limitations under the License.
23 */
24 
25 #pragma once
26 #include <cstdlib>
27 #include <cstdio>
28 #include <cstring>
29 #include <cstdarg>
30 
31 #include <string>
32 #include <iostream>
33 #include <unordered_map>
34 #include <cinttypes>
35 #include <string_view>
36 #include <array>
37 #include <netinet/in.h>
38 #include <arpa/inet.h>
39 
40 #include "ts/ts.h"
41 #include "ts/remap.h"
42 
43 typedef std::unordered_map<std::string, bool> OutstandingRequests;
44 const char PLUGIN_NAME[] = "cache_fill";
45 
46 class BgFetchState
47 {
48 public:
49   BgFetchState()                     = default;
50   BgFetchState(BgFetchState const &) = delete;
51   void operator=(BgFetchState const &) = delete;
52 
53   static BgFetchState &
getInstance()54   getInstance()
55   {
56     static BgFetchState _instance;
57     return _instance;
58   }
59 
~BgFetchState()60   ~BgFetchState() { TSMutexDestroy(_lock); }
61 
62   bool
acquire(const std::string & url)63   acquire(const std::string &url)
64   {
65     bool ret;
66 
67     TSMutexLock(_lock);
68     if (_urls.end() == _urls.find(url)) {
69       _urls[url] = true;
70       ret        = true;
71     } else {
72       ret = false;
73     }
74     TSMutexUnlock(_lock);
75 
76     TSDebug(PLUGIN_NAME, "BgFetchState.acquire(): ret = %d, url = %s", ret, url.c_str());
77 
78     return ret;
79   }
80 
81   bool
release(const std::string & url)82   release(const std::string &url)
83   {
84     bool ret;
85 
86     TSMutexLock(_lock);
87     if (_urls.end() == _urls.find(url)) {
88       ret = false;
89     } else {
90       _urls.erase(url);
91       ret = true;
92     }
93     TSMutexUnlock(_lock);
94 
95     return ret;
96   }
97 
98 private:
99   OutstandingRequests _urls;
100   TSMutex _lock = TSMutexCreate();
101 };
102 
103 //////////////////////////////////////////////////////////////////////////////
104 // Hold and manage some state for the TXN background fetch continuation.
105 // This is necessary, because the TXN is likely to not be available
106 // during the time we fetch from origin.
107 struct BgFetchData {
BgFetchDataBgFetchData108   BgFetchData() { memset(&client_ip, 0, sizeof(client_ip)); }
109 
~BgFetchDataBgFetchData110   ~BgFetchData()
111   {
112     TSHandleMLocRelease(mbuf, TS_NULL_MLOC, hdr_loc);
113     TSHandleMLocRelease(mbuf, TS_NULL_MLOC, url_loc);
114 
115     TSMBufferDestroy(mbuf);
116 
117     if (vc) {
118       TSError("[%s] Destroyed BgFetchDATA while VC was alive", PLUGIN_NAME);
119       TSVConnClose(vc);
120       vc = nullptr;
121     }
122 
123     // If we got schedule, also clean that up
124     if (_cont) {
125       releaseUrl();
126 
127       TSContDestroy(_cont);
128       _cont = nullptr;
129       TSIOBufferReaderFree(req_io_buf_reader);
130       TSIOBufferDestroy(req_io_buf);
131       TSIOBufferReaderFree(resp_io_buf_reader);
132       TSIOBufferDestroy(resp_io_buf);
133     }
134   }
135 
136   bool
acquireUrlBgFetchData137   acquireUrl() const
138   {
139     return BgFetchState::getInstance().acquire(_url);
140   }
141   bool
releaseUrlBgFetchData142   releaseUrl() const
143   {
144     return BgFetchState::getInstance().release(_url);
145   }
146 
147   const char *
getUrlBgFetchData148   getUrl() const
149   {
150     return _url.c_str();
151   }
152 
153   void
addBytesBgFetchData154   addBytes(int64_t b)
155   {
156     _bytes += b;
157   }
158 
159   bool initialize(TSMBuffer request, TSMLoc req_hdr, TSHttpTxn txnp);
160   void schedule();
161 
162   TSMBuffer mbuf = TSMBufferCreate();
163   TSMLoc hdr_loc = TS_NULL_MLOC;
164   TSMLoc url_loc = TS_NULL_MLOC;
165 
166   struct sockaddr_storage client_ip;
167 
168   // This is for the actual background fetch / NetVC
169   TSVConn vc                          = nullptr;
170   TSIOBuffer req_io_buf               = nullptr;
171   TSIOBuffer resp_io_buf              = nullptr;
172   TSIOBufferReader req_io_buf_reader  = nullptr;
173   TSIOBufferReader resp_io_buf_reader = nullptr;
174   TSVIO r_vio                         = nullptr;
175   TSVIO w_vio                         = nullptr;
176 
177 private:
178   std::string _url;
179   int64_t _bytes = 0;
180   TSCont _cont   = nullptr;
181 };
182