This is a fork of: https://github.com/roglew/puppy
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

226 lines
7.0 KiB

package main
import (
"errors"
"fmt"
)
type MessageStorage interface {
// NOTE: Load func responsible for loading dependent messages, delte function responsible for deleting dependent messages
// if it takes an ID, the storage is responsible for dependent messages
// Close the storage
Close()
// Update an existing request in the storage. Requires that it has already been saved
UpdateRequest(req *ProxyRequest) error
// Save a new instance of the request in the storage regardless of if it has already been saved
SaveNewRequest(req *ProxyRequest) error
// Load a request given a unique id
LoadRequest(reqid string) (*ProxyRequest, error)
LoadUnmangledRequest(reqid string) (*ProxyRequest, error)
// Delete a request
DeleteRequest(reqid string) error
// Update an existing response in the storage. Requires that it has already been saved
UpdateResponse(rsp *ProxyResponse) error
// Save a new instance of the response in the storage regardless of if it has already been saved
SaveNewResponse(rsp *ProxyResponse) error
// Load a response given a unique id
LoadResponse(rspid string) (*ProxyResponse, error)
LoadUnmangledResponse(rspid string) (*ProxyResponse, error)
// Delete a response
DeleteResponse(rspid string) error
// Update an existing websocket message in the storage. Requires that it has already been saved
UpdateWSMessage(req *ProxyRequest, wsm *ProxyWSMessage) error
// Save a new instance of the websocket message in the storage regardless of if it has already been saved
SaveNewWSMessage(req *ProxyRequest, wsm *ProxyWSMessage) error
// Load a websocket given a unique id
LoadWSMessage(wsmid string) (*ProxyWSMessage, error)
LoadUnmangledWSMessage(wsmid string) (*ProxyWSMessage, error)
// Delete a WSMessage
DeleteWSMessage(wsmid string) error
// Get list of all the request keys
RequestKeys() ([]string, error)
// A function to perform a search of requests in the storage. Same arguments as NewRequestChecker
Search(limit int64, args ...interface{}) ([]*ProxyRequest, error)
// A function to naively check every function in storage with the given function and return the ones that match
CheckRequests(limit int64, checker RequestChecker) ([]*ProxyRequest, error)
// Same as Search() but returns the IDs of the requests instead
// If Search() starts causing memory errors and I can't assume all the matching requests will fit in memory, I'll implement this or something
//SearchIDs(args ...interface{}) ([]string, error)
// Query functions
AllSavedQueries() ([]*SavedQuery, error)
SaveQuery(name string, query MessageQuery) error
LoadQuery(name string) (MessageQuery, error)
DeleteQuery(name string) error
}
const QueryNotSupported = ConstErr("custom query not supported")
type ReqSort []*ProxyRequest
type SavedQuery struct {
Name string
Query MessageQuery
}
func (reql ReqSort) Len() int {
return len(reql)
}
func (reql ReqSort) Swap(i int, j int) {
reql[i], reql[j] = reql[j], reql[i]
}
func (reql ReqSort) Less(i int, j int) bool {
return reql[j].StartDatetime.After(reql[i].StartDatetime)
}
type WSSort []*ProxyWSMessage
func (wsml WSSort) Len() int {
return len(wsml)
}
func (wsml WSSort) Swap(i int, j int) {
wsml[i], wsml[j] = wsml[j], wsml[i]
}
func (wsml WSSort) Less(i int, j int) bool {
return wsml[j].Timestamp.After(wsml[i].Timestamp)
}
/*
General storage functions
*/
func SaveNewRequest(ms MessageStorage, req *ProxyRequest) error {
if req.ServerResponse != nil {
if err := SaveNewResponse(ms, req.ServerResponse); err != nil {
return fmt.Errorf("error saving server response to request: %s", err.Error())
}
}
if req.Unmangled != nil {
if req.DbId != "" && req.DbId == req.Unmangled.DbId {
return errors.New("request has same DbId as unmangled version")
}
if err := SaveNewRequest(ms, req.Unmangled); err != nil {
return fmt.Errorf("error saving unmangled version of request: %s", err.Error())
}
}
if err := ms.SaveNewRequest(req); err != nil {
return fmt.Errorf("error saving new request: %s", err.Error())
}
for _, wsm := range req.WSMessages {
if err := SaveNewWSMessage(ms, req, wsm); err != nil {
return fmt.Errorf("error saving request's ws message: %s", err.Error())
}
}
return nil
}
func UpdateRequest(ms MessageStorage, req *ProxyRequest) error {
if req.ServerResponse != nil {
if err := UpdateResponse(ms, req.ServerResponse); err != nil {
return fmt.Errorf("error saving server response to request: %s", err.Error())
}
}
if req.Unmangled != nil {
if req.DbId != "" && req.DbId == req.Unmangled.DbId {
return errors.New("request has same DbId as unmangled version")
}
if err := UpdateRequest(ms, req.Unmangled); err != nil {
return fmt.Errorf("error saving unmangled version of request: %s", err.Error())
}
}
if req.DbId == "" {
if err := ms.SaveNewRequest(req); err != nil {
return fmt.Errorf("error saving new request: %s", err.Error())
}
} else {
if err := ms.UpdateRequest(req); err != nil {
return fmt.Errorf("error updating request: %s", err.Error())
}
}
for _, wsm := range req.WSMessages {
if err := UpdateWSMessage(ms, req, wsm); err != nil {
return fmt.Errorf("error saving request's ws message: %s", err.Error())
}
}
return nil
}
func SaveNewResponse(ms MessageStorage, rsp *ProxyResponse) error {
if rsp.Unmangled != nil {
if rsp.DbId != "" && rsp.DbId == rsp.Unmangled.DbId {
return errors.New("response has same DbId as unmangled version")
}
if err := SaveNewResponse(ms, rsp.Unmangled); err != nil {
return fmt.Errorf("error saving unmangled version of response: %s", err.Error())
}
}
return ms.SaveNewResponse(rsp)
}
func UpdateResponse(ms MessageStorage, rsp *ProxyResponse) error {
if rsp.Unmangled != nil {
if rsp.DbId != "" && rsp.DbId == rsp.Unmangled.DbId {
return errors.New("response has same DbId as unmangled version")
}
if err := UpdateResponse(ms, rsp.Unmangled); err != nil {
return fmt.Errorf("error saving unmangled version of response: %s", err.Error())
}
}
if rsp.DbId == "" {
return ms.SaveNewResponse(rsp)
} else {
return ms.UpdateResponse(rsp)
}
}
func SaveNewWSMessage(ms MessageStorage, req *ProxyRequest, wsm *ProxyWSMessage) error {
if wsm.Unmangled != nil {
if wsm.DbId != "" && wsm.DbId == wsm.Unmangled.DbId {
return errors.New("websocket message has same DbId as unmangled version")
}
if err := SaveNewWSMessage(ms, nil, wsm.Unmangled); err != nil {
return fmt.Errorf("error saving unmangled version of websocket message: %s", err.Error())
}
}
return ms.SaveNewWSMessage(req, wsm)
}
func UpdateWSMessage(ms MessageStorage, req *ProxyRequest, wsm *ProxyWSMessage) error {
if wsm.Unmangled != nil {
if wsm.DbId != "" && wsm.Unmangled.DbId == wsm.DbId {
return errors.New("websocket message has same DbId as unmangled version")
}
if err := UpdateWSMessage(ms, nil, wsm.Unmangled); err != nil {
return fmt.Errorf("error saving unmangled version of websocket message: %s", err.Error())
}
}
if wsm.DbId == "" {
return ms.SaveNewWSMessage(req, wsm)
} else {
return ms.UpdateWSMessage(req, wsm)
}
}