111 lines
2.6 KiB
Go
111 lines
2.6 KiB
Go
package main
|
|
|
|
import (
|
|
"net/http"
|
|
"net/http/httputil"
|
|
"net/url"
|
|
"strings"
|
|
"time"
|
|
|
|
"secure-ollama-api/utils"
|
|
|
|
"git.secnex.io/secnex/masterlog"
|
|
)
|
|
|
|
type Proxy struct {
|
|
ReverseProxy *httputil.ReverseProxy
|
|
TargetURL *url.URL
|
|
APIKey string
|
|
}
|
|
|
|
type responseWriter struct {
|
|
http.ResponseWriter
|
|
statusCode int
|
|
}
|
|
|
|
func NewProxy() (*Proxy, error) {
|
|
remote, err := url.Parse(utils.GetEnv("OLLAMA_SERVER_URL", "http://ollama:11434"))
|
|
if err != nil {
|
|
masterlog.Error(err.Error())
|
|
return nil, err
|
|
}
|
|
|
|
apiKey := utils.GetEnv("OLLAMA_API_KEY", "")
|
|
if apiKey == "" {
|
|
masterlog.Error("OLLAMA_API_KEY environment variable is not set")
|
|
}
|
|
|
|
return &Proxy{
|
|
ReverseProxy: httputil.NewSingleHostReverseProxy(remote),
|
|
TargetURL: remote,
|
|
APIKey: apiKey,
|
|
}, nil
|
|
}
|
|
|
|
func (rw *responseWriter) WriteHeader(code int) {
|
|
rw.statusCode = code
|
|
rw.ResponseWriter.WriteHeader(code)
|
|
}
|
|
|
|
func (p *Proxy) RequestHandler() func(w http.ResponseWriter, r *http.Request) {
|
|
return func(w http.ResponseWriter, r *http.Request) {
|
|
masterlog.Info("Request received", map[string]interface{}{
|
|
"url": r.URL.String(),
|
|
"ip": r.RemoteAddr,
|
|
})
|
|
startTime := time.Now()
|
|
|
|
// Check Authorization header
|
|
authHeader := r.Header.Get("Authorization")
|
|
expectedAuth := "Bearer " + p.APIKey
|
|
|
|
if p.APIKey == "" {
|
|
masterlog.Error("OLLAMA_API_KEY is not configured", map[string]interface{}{
|
|
"url": r.URL.String(),
|
|
"ip": r.RemoteAddr,
|
|
})
|
|
http.Error(w, "Internal server error", http.StatusInternalServerError)
|
|
return
|
|
}
|
|
|
|
if authHeader == "" || !strings.HasPrefix(authHeader, "Bearer ") {
|
|
masterlog.Info("Unauthorized request - missing or invalid Authorization header", map[string]interface{}{
|
|
"url": r.URL.String(),
|
|
"ip": r.RemoteAddr,
|
|
})
|
|
http.Error(w, "Unauthorized", http.StatusUnauthorized)
|
|
return
|
|
}
|
|
|
|
if authHeader != expectedAuth {
|
|
masterlog.Info("Unauthorized request - invalid API key", map[string]interface{}{
|
|
"url": r.URL.String(),
|
|
"ip": r.RemoteAddr,
|
|
})
|
|
http.Error(w, "Unauthorized", http.StatusUnauthorized)
|
|
return
|
|
}
|
|
|
|
// Wrap response writer to capture status code
|
|
rw := &responseWriter{
|
|
ResponseWriter: w,
|
|
statusCode: http.StatusOK, // default status code
|
|
}
|
|
|
|
p.ReverseProxy.ServeHTTP(rw, r)
|
|
|
|
duration := time.Since(startTime)
|
|
masterlog.Info("Request completed", map[string]interface{}{
|
|
"url": r.URL.String(),
|
|
"duration": duration.Milliseconds(),
|
|
"statusCode": rw.statusCode,
|
|
"ip": r.RemoteAddr,
|
|
})
|
|
}
|
|
}
|
|
|
|
func (p *Proxy) Start() {
|
|
http.HandleFunc("/", p.RequestHandler())
|
|
http.ListenAndServe(":8080", nil)
|
|
}
|