better task completion tracking, proxy tracking, and search page updates

MythicReactUI/src/components/MythicComponents/MythicTabPanel.js

@@ -172,10 +172,9 @@ export function MythicSearchTabLabel(props) {
     return (
         <Tab
             label={
-                <span>
-                    {label}
-                    <br />
+                <span className="mythic-search-tab-label">
                     {iconComponent}
+                    <span>{label}</span>
                 </span>
             }
             {...a11yProps(index)}

MythicReactUI/src/components/MythicComponents/MythicTableToolbar.js

@@ -17,9 +17,12 @@ export const MythicTableToolbar = ({children, className = "", style = {}, varian
     );
 };
 
-export const MythicTableToolbarGroup = ({children, grow = false, className = "", style = {}}) => {
+export const MythicTableToolbarGroup = ({children, grow = false, label, className = "", style = {}}) => {
     return (
         <Box className={`mythic-table-toolbar-group ${grow ? "mythic-table-toolbar-group-grow" : ""} ${className}`.trim()} style={style}>
+            {label &&
+                <span className="mythic-table-toolbar-group-label">{label}</span>
+            }
             {children}
         </Box>
     );

MythicReactUI/src/components/pages/Callbacks/CallbacksTabsTaskingInput.js

@@ -226,9 +226,11 @@ const getTaskingParameterTooltip = (parameter) => {
     const description = parameter.description ? `\n${parameter.description}` : "";
     return `${label} · ${parameter.parameter_type}${requiredText}${displayName}${internalName}${description}`;
 }
-const TaskingParameterPreviewChip = ({parameter, required=false, active=false}) => (
+const TaskingParameterPreviewChip = ({parameter, required=false, active=false, onClick}) => (
     <MythicStyledTooltip title={getTaskingParameterTooltip(parameter)}>
         <Chip
+            aria-label={onClick ? `Insert ${getTaskingParameterLabel(parameter)}` : getTaskingParameterLabel(parameter)}
+            clickable={Boolean(onClick)}
             className={`mythic-tasking-parameter-preview-chip${required ? " mythic-tasking-parameter-preview-chip-required" : ""}${active ? " mythic-tasking-parameter-preview-chip-active" : ""}`}
             label={
                 <span className="mythic-tasking-parameter-preview-chip-label">
@@ -239,6 +241,14 @@ const TaskingParameterPreviewChip = ({parameter, required=false, active=false})
                     <span className="mythic-tasking-parameter-preview-chip-type">{parameter.parameter_type}</span>
                 </span>
             }
+            onClick={onClick ? (event) => {
+                event.preventDefault();
+                event.stopPropagation();
+                onClick(parameter);
+            } : undefined}
+            onMouseDown={onClick ? (event) => {
+                event.preventDefault();
+            } : undefined}
             size="small"
         />
     </MythicStyledTooltip>
@@ -394,6 +404,7 @@ export function CallbacksTabsTaskingInputPreMemo(props){
     const toastId = "tasking-toast-message";
     const theme = useTheme();
     const inputRef = React.useRef(null);
+    const pendingCursorPosition = React.useRef(null);
     const snackMessageStyles = {position:"bottom-left", autoClose: 1000, toastId: toastId, style: {marginBottom: "30px"}};
     const snackReverseSearchMessageStyles = {position:"bottom-left", autoClose: 1000,  toastId: toastId, style: {marginBottom: "70px"}};
     const [commandPayloadType, setCommandPayloadType] = React.useState("");
@@ -641,6 +652,50 @@ export function CallbacksTabsTaskingInputPreMemo(props){
             setUnmodifiedHistoryValue("parsed_cli");
         }
     }
+    React.useEffect( () => {
+        if(pendingCursorPosition.current === null){
+            return;
+        }
+        const cursorPosition = pendingCursorPosition.current;
+        pendingCursorPosition.current = null;
+        window.requestAnimationFrame(() => {
+            if(inputRef.current){
+                inputRef.current.focus();
+                if(inputRef.current.setSelectionRange){
+                    const adjustedCursorPosition = Math.min(cursorPosition, inputRef.current.value.length);
+                    inputRef.current.setSelectionRange(adjustedCursorPosition, adjustedCursorPosition);
+                }
+            }
+        });
+    }, [message])
+    const insertParameterIntoCommandLine = (parameter) => {
+        if(!parameter?.cli_name){
+            return;
+        }
+        const input = inputRef.current;
+        let selectionStart = typeof input?.selectionStart === "number" ? input.selectionStart : message.length;
+        let selectionEnd = typeof input?.selectionEnd === "number" ? input.selectionEnd : selectionStart;
+        if(selectionStart === selectionEnd){
+            const beforeCursor = message.slice(0, selectionStart);
+            const partialParameterMatch = beforeCursor.match(/(^|\s)(-\S*)$/);
+            if(partialParameterMatch){
+                selectionStart -= partialParameterMatch[2].length;
+            }
+        }
+        const before = message.slice(0, selectionStart);
+        let after = message.slice(selectionEnd);
+        const parameterText = `-${parameter.cli_name}`;
+        const insertText = `${before.length > 0 && !/\s$/.test(before) ? " " : ""}${parameterText} `;
+        if(after.length > 0 && !/^\s/.test(after)){
+            after = ` ${after}`;
+        }
+        tabOptions.current = [];
+        tabOptionsType.current = "param_name";
+        tabOptionsIndex.current = 0;
+        setUnmodifiedHistoryValue("parsed_cli");
+        pendingCursorPosition.current = before.length + insertText.length;
+        setMessage(`${before}${insertText}${after}`);
+    }
     const onKeyDown = (event) => {
         if(event.key === "Enter" && (event.ctrlKey || event.metaKey)){
             setMessage(message + "\n");
@@ -2117,7 +2172,7 @@ export function CallbacksTabsTaskingInputPreMemo(props){
                         <TaskingParameterPreviewChip key={"active" + commandParameterPreview.activeParameter.id} parameter={commandParameterPreview.activeParameter} active={true} />
                     }
                     {commandParameterPreview.requiredParameters.slice(0, commandParameterPreview.activeParameter ? 5 : 6).map((parameter) => (
-                        <TaskingParameterPreviewChip key={"required" + parameter.id} parameter={parameter} required={true} />
+                        <TaskingParameterPreviewChip key={"required" + parameter.id} parameter={parameter} required={true} onClick={insertParameterIntoCommandLine} />
                     ))}
                     {commandParameterPreview.requiredParameters.length > (commandParameterPreview.activeParameter ? 5 : 6) &&
                         <span className="mythic-tasking-parameter-preview-more">
@@ -2125,7 +2180,7 @@ export function CallbacksTabsTaskingInputPreMemo(props){
                         </span>
                     }
                     {commandParameterPreview.optionalParameters.slice(0, commandParameterPreview.requiredParameters.length > 0 ? 4 : (commandParameterPreview.activeParameter ? 5 : 6)).map((parameter) => (
-                        <TaskingParameterPreviewChip key={"optional" + parameter.id} parameter={parameter} />
+                        <TaskingParameterPreviewChip key={"optional" + parameter.id} parameter={parameter} onClick={insertParameterIntoCommandLine} />
                     ))}
                     {commandParameterPreview.optionalParameters.length > (commandParameterPreview.requiredParameters.length > 0 ? 4 : (commandParameterPreview.activeParameter ? 5 : 6)) &&
                         <span className="mythic-tasking-parameter-preview-more">

MythicReactUI/src/components/pages/Search/Search.js

@@ -97,7 +97,7 @@ export function Search(props){
           <AppBar
             position="static"
             color="default"
-            className={"no-box-shadow"}
+            className={"no-box-shadow mythic-search-tabs-bar"}
             sx={(theme) => ({
               backgroundColor: theme.surfaces?.muted || theme.palette.background.paper,
               border: `1px solid ${theme.table?.borderSoft || theme.borderColor}`,
@@ -106,6 +106,7 @@ export function Search(props){
             })}
           >
             <Tabs
+              className="mythic-search-tabs"
               value={value}
               onChange={handleChange}
               indicatorColor="primary"

MythicReactUI/src/components/pages/Search/SearchTabArtifacts.js

@@ -233,11 +233,11 @@ const SearchTabArtifactsSearchPanel = (props) => {
         }
     }, [props.value, props.index])
     return (
-        <MythicTableToolbar>
-            <MythicTableToolbarGroup grow>
+        <MythicTableToolbar variant="search">
+            <MythicTableToolbarGroup grow label="Search">
                 <MythicSearchField value={search} onChange={handleSearchValueChange} onEnter={submitSearch} onSearch={submitSearch} />
             </MythicTableToolbarGroup>
-            <MythicTableToolbarGroup>
+            <MythicTableToolbarGroup label="In">
                 <MythicToolbarSelect
                     value={searchField}
                     onChange={handleSearchFieldChange}
@@ -248,6 +248,8 @@ const SearchTabArtifactsSearchPanel = (props) => {
                         ))
                     }
                 </MythicToolbarSelect>
+            </MythicTableToolbarGroup>
+            <MythicTableToolbarGroup label="Cleanup">
                 <MythicToolbarSelect
                     value={cleanupField}
                     onChange={handleCleanupChange}
@@ -259,7 +261,7 @@ const SearchTabArtifactsSearchPanel = (props) => {
                     }
                 </MythicToolbarSelect>
             </MythicTableToolbarGroup>
-            <MythicTableToolbarGroup>
+            <MythicTableToolbarGroup label="Actions">
                 {createArtifactDialogOpen &&
                     <MythicDialog fullWidth={true} maxWidth="md" open={createArtifactDialogOpen}
                                   onClose={()=>{setCreateArtifactDialogOpen(false);}}

MythicReactUI/src/components/pages/Search/SearchTabCallbacks.js

@@ -231,10 +231,10 @@ const SearchTabCallbacksSearchPanel = (props) => {
     }, [props.value, props.index])
     return (
         <MythicTableToolbar variant="search">
-            <MythicTableToolbarGroup grow>
+            <MythicTableToolbarGroup grow label="Search">
                 <MythicSearchField value={search} onChange={handleSearchValueChange} onEnter={submitSearch} onSearch={submitSearch} />
             </MythicTableToolbarGroup>
-            <MythicTableToolbarGroup>
+            <MythicTableToolbarGroup label="In">
                 <MythicToolbarSelect
                     value={searchField}
                     onChange={handleSearchFieldChange}

MythicReactUI/src/components/pages/Search/SearchTabCredentials.js

@@ -217,10 +217,10 @@ const SearchTabCredentialsSearchPanel = (props) => {
     }, [props.value, props.index])
     return (
         <MythicTableToolbar variant="search">
-            <MythicTableToolbarGroup grow>
+            <MythicTableToolbarGroup grow label="Search">
                 <MythicSearchField value={search} onChange={handleSearchValueChange} onEnter={submitSearch} onSearch={submitSearch} />
             </MythicTableToolbarGroup>
-            <MythicTableToolbarGroup>
+            <MythicTableToolbarGroup label="In">
                 <MythicToolbarSelect
                     value={searchField}
                     onChange={handleSearchFieldChange}
@@ -232,7 +232,7 @@ const SearchTabCredentialsSearchPanel = (props) => {
                     }
                 </MythicToolbarSelect>
             </MythicTableToolbarGroup>
-            <MythicTableToolbarGroup>
+            <MythicTableToolbarGroup label="Actions">
                 {createCredentialDialogOpen &&
                     <MythicDialog fullWidth={true} maxWidth="md" open={createCredentialDialogOpen} 
                         onClose={()=>{setCreateCredentialDialogOpen(false);}} 

MythicReactUI/src/components/pages/Search/SearchTabCustomBrowsers.js

@@ -195,8 +195,8 @@ const SearchTabCustomBrowserSearchPanel = (props) => {
         }
     }, [props.value, props.index]);
     return (
-        <MythicTableToolbar>
-            <MythicTableToolbarGroup>
+        <MythicTableToolbar variant="search">
+            <MythicTableToolbarGroup label="Browser">
                 <MythicToolbarSelect
                     value={selectedBrowser}
                     onChange={handleBrowserChange}
@@ -208,14 +208,14 @@ const SearchTabCustomBrowserSearchPanel = (props) => {
                     }
                 </MythicToolbarSelect>
             </MythicTableToolbarGroup>
-            <MythicTableToolbarGroup style={{minWidth: "13rem"}}>
+            <MythicTableToolbarGroup label="Host" style={{minWidth: "13rem"}}>
                 <MythicSearchField placeholder="Host Name Search..." name="Host" value={searchHost}
                                    onChange={handleSearchHostValueChange} onEnter={submitSearch}/>
             </MythicTableToolbarGroup>
-            <MythicTableToolbarGroup grow>
+            <MythicTableToolbarGroup grow label="Search">
                 <MythicSearchField value={search} onChange={handleSearchValueChange} onEnter={submitSearch} onSearch={submitSearch} />
             </MythicTableToolbarGroup>
-            <MythicTableToolbarGroup>
+            <MythicTableToolbarGroup label="In">
                 <MythicToolbarSelect
                     value={searchField}
                     onChange={handleSearchFieldChange}

MythicReactUI/src/components/pages/Search/SearchTabFiles.js

@@ -553,14 +553,14 @@ const SearchTabFilesSearchPanel = (props) => {
     }, [props.value, props.index]);
     return (
         <MythicTableToolbar variant="search">
-            <MythicTableToolbarGroup style={{minWidth: "13rem"}}>
+            <MythicTableToolbarGroup label="Host" style={{minWidth: "13rem"}}>
                 <MythicSearchField placeholder="Host Name Search..." name="Host" value={searchHost}
                                    onChange={handleSearchHostValueChange} onEnter={submitSearch}/>
             </MythicTableToolbarGroup>
-            <MythicTableToolbarGroup grow>
+            <MythicTableToolbarGroup grow label="Search">
                 <MythicSearchField value={search} onChange={handleSearchValueChange} onEnter={submitSearch} onSearch={submitSearch} />
             </MythicTableToolbarGroup>
-            <MythicTableToolbarGroup>
+            <MythicTableToolbarGroup label="In">
                 <MythicToolbarSelect
                     value={searchField}
                     onChange={handleSearchFieldChange}
@@ -572,7 +572,7 @@ const SearchTabFilesSearchPanel = (props) => {
                     }
                 </MythicToolbarSelect>
             </MythicTableToolbarGroup>
-            <MythicTableToolbarGroup>
+            <MythicTableToolbarGroup label="Location">
                 <MythicToolbarSelect
                     value={searchLocation}
                     onChange={handleSearchLocationChange}
@@ -584,7 +584,7 @@ const SearchTabFilesSearchPanel = (props) => {
                     }
                 </MythicToolbarSelect>
             </MythicTableToolbarGroup>
-            <MythicTableToolbarGroup>
+            <MythicTableToolbarGroup label="Actions">
                 <MythicToolbarButton className="mythic-toolbar-button-hover-success" variant="outlined" component="label" startIcon={<BackupIcon />}>
                     Files
                     <input onChange={onFileChange} type="file" multiple hidden/>

MythicReactUI/src/components/pages/Search/SearchTabKeylog.js

@@ -172,11 +172,11 @@ const SearchTabKeylogsSearchPanel = (props) => {
         }
     }, [props.value, props.index])
     return (
-        <MythicTableToolbar>
-            <MythicTableToolbarGroup grow>
+        <MythicTableToolbar variant="search">
+            <MythicTableToolbarGroup grow label="Search">
                 <MythicSearchField value={search} onChange={handleSearchValueChange} onEnter={submitSearch} onSearch={submitSearch} />
             </MythicTableToolbarGroup>
-            <MythicTableToolbarGroup>
+            <MythicTableToolbarGroup label="In">
                 <MythicToolbarSelect
                     value={searchField}
                     onChange={handleSearchFieldChange}

MythicReactUI/src/components/pages/Search/SearchTabPayloads.js

@@ -232,10 +232,10 @@ const SearchTabPayloadsSearchPanel = (props) => {
     }, []);
     return (
         <MythicTableToolbar variant="search">
-            <MythicTableToolbarGroup grow>
+            <MythicTableToolbarGroup grow label="Search">
                 <MythicSearchField value={search} onChange={handleSearchValueChange} onEnter={submitSearch} onSearch={submitSearch} />
             </MythicTableToolbarGroup>
-            <MythicTableToolbarGroup>
+            <MythicTableToolbarGroup label="In">
                 <MythicToolbarSelect
                     value={searchField}
                     onChange={handleSearchFieldChange}
@@ -247,7 +247,7 @@ const SearchTabPayloadsSearchPanel = (props) => {
                     }
                 </MythicToolbarSelect>
             </MythicTableToolbarGroup>
-            <MythicTableToolbarGroup>
+            <MythicTableToolbarGroup label="C2 profile">
                 <MythicToolbarSelect
                     value={searchC2}
                     onChange={handleC2FieldChange}
@@ -260,7 +260,7 @@ const SearchTabPayloadsSearchPanel = (props) => {
                 </MythicToolbarSelect>
 
             </MythicTableToolbarGroup>
-            <MythicTableToolbarGroup>
+            <MythicTableToolbarGroup label="Payload type">
                 <MythicToolbarSelect
                     value={searchPayloadType}
                     onChange={handlePayloadTypeFieldChange}
@@ -272,7 +272,7 @@ const SearchTabPayloadsSearchPanel = (props) => {
                     }
                 </MythicToolbarSelect>
             </MythicTableToolbarGroup>
-            <MythicTableToolbarGroup>
+            <MythicTableToolbarGroup label="Filters">
                 <MythicToolbarToggle
                     checked={showDeleted}
                     onClick={handleToggleShowDeleted}

MythicReactUI/src/components/pages/Search/SearchTabProcesses.js

@@ -153,15 +153,15 @@ const SearchTabProcessesSearchPanel = (props) => {
         }
     }, [props.value, props.index]);
     return (
-        <MythicTableToolbar>
-            <MythicTableToolbarGroup style={{minWidth: "13rem"}}>
+        <MythicTableToolbar variant="search">
+            <MythicTableToolbarGroup label="Host" style={{minWidth: "13rem"}}>
                 <MythicSearchField placeholder="Host Name Search..." name="Host" value={searchHost}
                                    onChange={handleSearchHostValueChange} onEnter={submitSearch}/>
             </MythicTableToolbarGroup>
-            <MythicTableToolbarGroup grow>
+            <MythicTableToolbarGroup grow label="Search">
                 <MythicSearchField value={search} onChange={handleSearchValueChange} onEnter={submitSearch} onSearch={submitSearch} />
             </MythicTableToolbarGroup>
-            <MythicTableToolbarGroup>
+            <MythicTableToolbarGroup label="In">
                 <MythicToolbarSelect
                     value={searchField}
                     onChange={handleSearchFieldChange}

MythicReactUI/src/components/pages/Search/SearchTabProxies.js

@@ -7,7 +7,7 @@ import {FontAwesomeIcon} from '@fortawesome/react-fontawesome';
 import {faSocks} from '@fortawesome/free-solid-svg-icons';
 import VisibilityOffIcon from '@mui/icons-material/VisibilityOff';
 import VisibilityIcon from '@mui/icons-material/Visibility';
-import {MythicTableToolbar, MythicTableToolbarGroup, MythicToolbarToggle} from "../../MythicComponents/MythicTableToolbar";
+import {MythicSearchField, MythicTableToolbar, MythicTableToolbarGroup, MythicToolbarToggle} from "../../MythicComponents/MythicTableToolbar";
 import {MythicSearchEmptyState} from "../../MythicComponents/MythicStateDisplay";
 
 const callbackPortsSub = gql`
@@ -55,6 +55,29 @@ export function SearchTabSocksLabel(props){
 export const SearchTabSocksPanel = (props) =>{
     const [callbackData, setCallbackData] = React.useState([]);
     const [showDeleted, setShowDeleted] = React.useState(false);
+    const [search, setSearch] = React.useState("");
+    const handleSearchValueChange = (name, value, error) => {
+        setSearch(value);
+    }
+    const submitSearch = () => {
+        props.changeSearchParam("search", search);
+    }
+    const handleToggleShowDeleted = () => {
+        const nextShowDeleted = !showDeleted;
+        setShowDeleted(nextShowDeleted);
+        props.changeSearchParam("showStopped", nextShowDeleted ? "true" : "false");
+    }
+    React.useEffect(() => {
+        if(props.value === props.index){
+            const queryParams = new URLSearchParams(window.location.search);
+            if(queryParams.has("search")){
+                setSearch(queryParams.get("search"));
+            }
+            if(queryParams.has("showStopped")){
+                setShowDeleted(queryParams.get("showStopped") === "true");
+            }
+        }
+    }, [props.value, props.index]);
     useSubscription(callbackPortsSub, {
         fetchPolicy: "no-cache",
         onData: ({data}) => {
@@ -79,14 +102,48 @@ export const SearchTabSocksPanel = (props) =>{
             console.log(data);
         }
     });
+    const filteredCallbacks = React.useMemo(() => {
+        const normalizedSearch = search.trim().toLowerCase();
+        if(normalizedSearch === ""){
+            return callbackData;
+        }
+        return callbackData.filter((proxy) => {
+            const searchableText = [
+                proxy.callback?.user,
+                proxy.callback?.host,
+                proxy.callback?.description,
+                proxy.callback?.domain,
+                proxy.callback?.display_id ? `C-${proxy.callback.display_id}` : "",
+                proxy.task?.display_id ? `T-${proxy.task.display_id}` : "",
+                proxy.local_port,
+                proxy.remote_ip,
+                proxy.remote_port,
+                proxy.port_type,
+                proxy.username,
+            ].join(" ").toLowerCase();
+            return searchableText.includes(normalizedSearch);
+        });
+    }, [callbackData, search]);
+    const visibleCallbacks = React.useMemo(() => {
+        return filteredCallbacks.filter((proxy) => showDeleted || !proxy.deleted);
+    }, [filteredCallbacks, showDeleted]);
 
     return (
         <MythicTabPanel {...props} >
-            <MythicTableToolbar>
-                <MythicTableToolbarGroup grow>
+            <MythicTableToolbar variant="search">
+                <MythicTableToolbarGroup grow label="Search">
+                    <MythicSearchField
+                        value={search}
+                        onChange={handleSearchValueChange}
+                        onEnter={submitSearch}
+                        onSearch={submitSearch}
+                        placeholder="Search proxies..."
+                    />
+                </MythicTableToolbarGroup>
+                <MythicTableToolbarGroup label="Filters">
                     <MythicToolbarToggle
                         checked={showDeleted}
-                        onClick={() => setShowDeleted(!showDeleted)}
+                        onClick={handleToggleShowDeleted}
                         label="Stopped"
                         activeIcon={<VisibilityIcon fontSize="small" />}
                         inactiveIcon={<VisibilityOffIcon fontSize="small" />}
@@ -94,11 +151,13 @@ export const SearchTabSocksPanel = (props) =>{
                 </MythicTableToolbarGroup>
             </MythicTableToolbar>
             <div style={{overflowY: "auto", height: "100%", display: "flex", flexDirection: "column"}}>
-                {callbackData.length > 0 ? (
-                    <ProxySearchTable callbacks={callbackData} showDeleted={showDeleted} />) : (
+                {visibleCallbacks.length > 0 ? (
+                    <ProxySearchTable callbacks={filteredCallbacks} showDeleted={showDeleted} />) : (
                     <MythicSearchEmptyState
                         compact
-                        description="No callback port forwards match the current stopped filter."
+                        description={search.trim() === "" ?
+                            "No callback port forwards match the current stopped filter." :
+                            "No callback port forwards match the current search and stopped filter."}
                         minHeight={180}
                     />
                 )}

MythicReactUI/src/components/pages/Search/SearchTabTags.js

@@ -195,11 +195,11 @@ const SearchTabTagsSearchPanel = (props) => {
         }
     }, [props.value, props.index]);
     return (
-        <MythicTableToolbar>
-            <MythicTableToolbarGroup grow>
+        <MythicTableToolbar variant="search">
+            <MythicTableToolbarGroup grow label="Search">
                 <MythicSearchField value={search} onChange={handleSearchValueChange} onEnter={submitSearch} onSearch={submitSearch} />
             </MythicTableToolbarGroup>
-            <MythicTableToolbarGroup>
+            <MythicTableToolbarGroup label="In">
                 <MythicToolbarSelect
                     value={searchField}
                     onChange={handleSearchFieldChange}

MythicReactUI/src/components/pages/Search/SearchTabTasks.js

@@ -257,15 +257,15 @@ const SearchTabTasksSearchPanel = (props) => {
     }, [props.value, props.index]);
     return (
         <MythicTableToolbar variant="search">
-            <MythicTableToolbarGroup style={{minWidth: "12rem"}}>
+            <MythicTableToolbarGroup label="Host" style={{minWidth: "12rem"}}>
                 <MythicSearchField disabled={props.alreadySearching} placeholder="Host..." name="Host" value={host}
                                    onChange={handleHostValueChange} onEnter={submitSearch}/>
             </MythicTableToolbarGroup>
-            <MythicTableToolbarGroup grow>
+            <MythicTableToolbarGroup grow label="Search">
                 <MythicSearchField disabled={props.alreadySearching} value={search}
                                    onChange={handleSearchValueChange} onEnter={submitSearch} onSearch={submitSearch} />
             </MythicTableToolbarGroup>
-            <MythicTableToolbarGroup>
+            <MythicTableToolbarGroup label="In">
                 <MythicToolbarSelect
                     value={searchField}
                     disabled={props.alreadySearching}
@@ -278,12 +278,12 @@ const SearchTabTasksSearchPanel = (props) => {
                     }
                 </MythicToolbarSelect>
             </MythicTableToolbarGroup>
-            <MythicTableToolbarGroup style={{minWidth: "13rem"}}>
+            <MythicTableToolbarGroup label="Status" style={{minWidth: "13rem"}}>
                 <MythicSearchField disabled={props.alreadySearching} placeholder="Filter Task Status..." name="Status"
                                    value={filterTaskStatus} onChange={handleFilterTaskStatusValueChange}
                                    onEnter={submitSearch}/>
             </MythicTableToolbarGroup>
-            <MythicTableToolbarGroup>
+            <MythicTableToolbarGroup label="Operator">
                 <MythicToolbarSelect
                     value={filterOperator}
                     disabled={props.alreadySearching}

MythicReactUI/src/components/pages/Search/SearchTabTokens.js

@@ -131,11 +131,11 @@ const SearchTabTokensSearchPanel = (props) => {
         }
     }, [props.value, props.index])
     return (
-        <MythicTableToolbar>
-            <MythicTableToolbarGroup grow>
+        <MythicTableToolbar variant="search">
+            <MythicTableToolbarGroup grow label="Search">
                 <MythicSearchField value={search} onChange={handleSearchValueChange} onEnter={submitSearch} onSearch={submitSearch} />
             </MythicTableToolbarGroup>
-            <MythicTableToolbarGroup>
+            <MythicTableToolbarGroup label="In">
                 <MythicToolbarSelect
                     value={searchField}
                     onChange={handleSearchFieldChange}

MythicReactUI/src/themes/GlobalStyles.js

@@ -585,11 +585,21 @@ tspan {
     height: 22px;
     max-width: min(16rem, 100%);
 }
+.mythic-tasking-parameter-preview-chip.MuiChip-clickable:hover {
+    background-color: ${(props) => alpha(props.theme.palette.info.main, props.theme.palette.mode === "dark" ? 0.12 : 0.07)};
+    border-color: ${(props) => alpha(props.theme.palette.info.main, props.theme.palette.mode === "dark" ? 0.36 : 0.24)};
+    color: ${(props) => props.theme.palette.info.main};
+}
 .mythic-tasking-parameter-preview-chip-required.MuiChip-root {
     background-color: ${(props) => alpha(props.theme.palette.warning.main, props.theme.palette.mode === "dark" ? 0.14 : 0.08)};
     border-color: ${(props) => alpha(props.theme.palette.warning.main, props.theme.palette.mode === "dark" ? 0.36 : 0.24)};
     color: ${(props) => props.theme.palette.warning.main};
 }
+.mythic-tasking-parameter-preview-chip-required.MuiChip-clickable:hover {
+    background-color: ${(props) => alpha(props.theme.palette.warning.main, props.theme.palette.mode === "dark" ? 0.18 : 0.1)};
+    border-color: ${(props) => alpha(props.theme.palette.warning.main, props.theme.palette.mode === "dark" ? 0.48 : 0.32)};
+    color: ${(props) => props.theme.palette.warning.main};
+}
 .mythic-tasking-parameter-preview-chip-active.MuiChip-root {
     background-color: ${(props) => alpha(props.theme.palette.info.main, props.theme.palette.mode === "dark" ? 0.14 : 0.08)};
     border-color: ${(props) => alpha(props.theme.palette.info.main, props.theme.palette.mode === "dark" ? 0.42 : 0.28)};
@@ -2627,10 +2637,11 @@ tspan {
 .mythic-table-toolbar-search {
     align-items: stretch;
     background-color: ${(props) => props.theme.palette.mode === "dark" ? "rgba(255,255,255,0.028)" : "rgba(0,0,0,0.012)"};
-    gap: 0.45rem;
-    padding: 0.45rem;
+    gap: 0.55rem;
+    padding: 0.55rem;
 }
 .mythic-table-toolbar-group {
+    align-content: center;
     align-items: center;
     display: flex;
     flex: 0 1 auto;
@@ -2642,6 +2653,14 @@ tspan {
 .mythic-table-toolbar-search .mythic-table-toolbar-group {
     gap: 0.35rem;
 }
+.mythic-table-toolbar-group-label {
+    color: ${(props) => props.theme.palette.text.secondary};
+    flex: 0 0 100%;
+    font-size: 0.66rem;
+    font-weight: 800;
+    line-height: 1;
+    padding-left: 0.1rem;
+}
 .mythic-table-toolbar-group-grow {
     flex: 1 1 20rem;
     min-width: min(100%, 16rem);
@@ -8443,6 +8462,30 @@ tspan {
 .MuiTabs-scrollButtons.Mui-disabled {
   opacity: 0.3;
 }
+.mythic-search-tabs-bar.MuiPaper-root {
+    background-image: ${getSectionHeaderGradient};
+}
+.mythic-search-tabs.MuiTabs-root {
+    min-height: 42px;
+}
+.mythic-search-tabs .MuiTabs-flexContainer {
+    flex-wrap: nowrap;
+}
+.mythic-search-tabs .MuiTab-root {
+    min-height: 42px;
+    padding: 0.45rem 0.7rem;
+}
+.mythic-search-tab-label {
+    align-items: center;
+    display: inline-flex;
+    gap: 0.42rem;
+    white-space: nowrap;
+}
+.mythic-search-tab-label .MuiSvgIcon-root,
+.mythic-search-tab-label .svg-inline--fa {
+    color: inherit;
+    font-size: 1rem;
+}
 .ace_editor{
     background-color: ${(props) => props.theme.outputBackgroundColor + "20"};
 }

mythic-docker/src/database/migrations/003003020_3.3.20.sql

@@ -76,6 +76,11 @@ create unique index if not exists response_task_sequence_number_unique
 on "public"."response" using btree (task_id, sequence_number)
 where sequence_number is not null;
 
+alter table "public"."task"
+    add column if not exists subtask_callback_function_started boolean not null default false,
+    add column if not exists group_callback_function_started boolean not null default false,
+    add column if not exists completed_callback_function_started boolean not null default false;
+
 -- +migrate StatementBegin
 create or replace function public.new_task_display_id() returns trigger
     language plpgsql
@@ -201,4 +206,8 @@ $$;
 drop index if exists "public"."callback_operation_display_id_unique";
 drop index if exists "public"."task_operation_display_id_unique";
 drop index if exists "public"."response_task_sequence_number_unique";
+alter table "public"."task"
+    drop column if exists completed_callback_function_started,
+    drop column if exists group_callback_function_started,
+    drop column if exists subtask_callback_function_started;
 drop table if exists "public"."operation_display_counters";

mythic-docker/src/database/structs/Task.go

@@ -47,10 +47,13 @@ type Task struct {
 	ParentTaskID                       structs.NullInt64 `db:"parent_task_id" json:"parent_task_id"`
 	SubtaskCallbackFunction            string            `db:"subtask_callback_function" json:"subtask_callback_function"`
 	SubtaskCallbackFunctionCompleted   bool              `db:"subtask_callback_function_completed" json:"subtask_callback_function_completed"`
+	SubtaskCallbackFunctionStarted     bool              `db:"subtask_callback_function_started" json:"subtask_callback_function_started"`
 	GroupCallbackFunction              string            `db:"group_callback_function" json:"group_callback_function"`
 	GroupCallbackFunctionCompleted     bool              `db:"group_callback_function_completed" json:"group_callback_function_completed"`
+	GroupCallbackFunctionStarted       bool              `db:"group_callback_function_started" json:"group_callback_function_started"`
 	CompletedCallbackFunction          string            `db:"completed_callback_function" json:"completed_callback_function"`
 	CompletedCallbackFunctionCompleted bool              `db:"completed_callback_function_completed" json:"completed_callback_function_completed"`
+	CompletedCallbackFunctionStarted   bool              `db:"completed_callback_function_started" json:"completed_callback_function_started"`
 	SubtaskGroupName                   string            `db:"subtask_group_name" json:"subtask_group_name"`
 	TaskingLocation                    string            `db:"tasking_location" json:"tasking_location"`
 	ParameterGroupName                 string            `db:"parameter_group_name" json:"parameter_group_name"`

mythic-docker/src/rabbitmq/util_agent_message_actions_post_response.go

@@ -294,14 +294,14 @@ type writeDownloadChunkToDisk struct {
 
 var writeDownloadChunkToDiskChan = make(chan writeDownloadChunkToDisk)
 
-type agentAgentMessagePostResponseChannelMessage struct {
+type agentMessagePostResponseUserOutputChannelMessage struct {
 	Response    string
 	SequenceNum *int64
 	Task        databaseStructs.Task
 }
 
-var asyncAgentMessagePostResponseWorkerChannels []chan agentAgentMessagePostResponseChannelMessage
-var asyncAgentMessagePostResponseWorkersOnce sync.Once
+var asyncAgentMessagePostResponseUserOutputWorkerChannels []chan agentMessagePostResponseUserOutputChannelMessage
+var asyncAgentMessagePostResponseUserOutputWorkersOnce sync.Once
 
 var responseInterceptMapOperationIDToEventGroupID = make(map[int]int)
 var responseInterceptMapLock = sync.RWMutex{}
@@ -320,25 +320,25 @@ func getAgentMessagePostResponseQueueSize() int {
 	return defaultAgentMessagePostResponseQueueSize
 }
 
-func initializeAsyncAgentMessagePostResponseWorkers() {
+func initializeAsyncAgentMessagePostResponseUserOutputWorkers() {
 	UpdateCachedResponseIntercept()
 	workerCount := getAgentMessagePostResponseWorkerCount()
 	queueSize := getAgentMessagePostResponseQueueSize()
-	asyncAgentMessagePostResponseWorkerChannels = make([]chan agentAgentMessagePostResponseChannelMessage, workerCount)
+	asyncAgentMessagePostResponseUserOutputWorkerChannels = make([]chan agentMessagePostResponseUserOutputChannelMessage, workerCount)
 	for i := 0; i < workerCount; i++ {
-		asyncAgentMessagePostResponseWorkerChannels[i] = make(chan agentAgentMessagePostResponseChannelMessage, queueSize)
-		go listenForAsyncAgentMessagePostResponseWorker(i, asyncAgentMessagePostResponseWorkerChannels[i])
+		asyncAgentMessagePostResponseUserOutputWorkerChannels[i] = make(chan agentMessagePostResponseUserOutputChannelMessage, queueSize)
+		go listenForAsyncAgentMessagePostResponseUserOutputWorker(i, asyncAgentMessagePostResponseUserOutputWorkerChannels[i])
 	}
 }
 
-func enqueueAsyncAgentMessagePostResponse(msg agentAgentMessagePostResponseChannelMessage) {
-	asyncAgentMessagePostResponseWorkersOnce.Do(initializeAsyncAgentMessagePostResponseWorkers)
-	workerCount := len(asyncAgentMessagePostResponseWorkerChannels)
+func enqueueAsyncAgentMessagePostResponse(msg agentMessagePostResponseUserOutputChannelMessage) {
+	asyncAgentMessagePostResponseUserOutputWorkersOnce.Do(initializeAsyncAgentMessagePostResponseUserOutputWorkers)
+	workerCount := len(asyncAgentMessagePostResponseUserOutputWorkerChannels)
 	workerID := msg.Task.ID % workerCount
 	if workerID < 0 {
 		workerID = 0
 	}
-	asyncAgentMessagePostResponseWorkerChannels[workerID] <- msg
+	asyncAgentMessagePostResponseUserOutputWorkerChannels[workerID] <- msg
 }
 
 func UpdateCachedResponseIntercept() {
@@ -374,10 +374,10 @@ func getCachedResponseInterceptEventGroupID(operationID int) (int, bool) {
 }
 
 func listenForAsyncAgentMessagePostResponseContent() {
-	asyncAgentMessagePostResponseWorkersOnce.Do(initializeAsyncAgentMessagePostResponseWorkers)
+	asyncAgentMessagePostResponseUserOutputWorkersOnce.Do(initializeAsyncAgentMessagePostResponseUserOutputWorkers)
 }
 
-func listenForAsyncAgentMessagePostResponseWorker(workerID int, input <-chan agentAgentMessagePostResponseChannelMessage) {
+func listenForAsyncAgentMessagePostResponseUserOutputWorker(workerID int, input <-chan agentMessagePostResponseUserOutputChannelMessage) {
 	insertStatement, err := database.DB.Preparex(insertAgentMessagePostResponseUserOutputQuery)
 	if err != nil {
 		logging.LogError(err, "Failed to prepare user_output insert statement for post_response worker", "worker_id", workerID)
@@ -386,11 +386,11 @@ func listenForAsyncAgentMessagePostResponseWorker(workerID int, input <-chan age
 		defer insertStatement.Close()
 	}
 	for msg := range input {
-		processAsyncAgentMessagePostResponseContent(msg, insertStatement)
+		processAsyncAgentMessagePostResponseUserOutput(msg, insertStatement)
 	}
 }
 
-func processAsyncAgentMessagePostResponseContent(msg agentAgentMessagePostResponseChannelMessage, insertStatement *sqlx.Stmt) {
+func processAsyncAgentMessagePostResponseUserOutput(msg agentMessagePostResponseUserOutputChannelMessage, insertStatement *sqlx.Stmt) {
 	// force chunking user_output can be useful, but might also affect command's browserscripts
 	//chunkSize := 1024 * 1024
 	//chunks := int(len(msg.Response)/chunkSize) + 1 // 1MB chunks
@@ -451,6 +451,44 @@ func getAgentMessagePostResponseTasks(responses []agentMessagePostResponse) (map
 	return tasksByAgentTaskID, nil
 }
 
+func updateAgentMessagePostResponseTask(task databaseStructs.Task) (bool, error) {
+	tx, err := database.DB.Beginx()
+	if err != nil {
+		return false, err
+	}
+	defer tx.Rollback()
+
+	wasCompleted := false
+	if err = tx.Get(&wasCompleted, `SELECT completed FROM task WHERE id=$1 FOR UPDATE`, task.ID); err != nil {
+		return false, err
+	}
+
+	updatedCompleted := false
+	err = tx.Get(&updatedCompleted, `UPDATE task SET
+		status=CASE WHEN completed THEN status ELSE $2 END,
+		completed=(completed OR $3),
+		status_timestamp_processed=coalesce(status_timestamp_processed, $4),
+		"timestamp"=$5,
+		stdout=$6,
+		stderr=$7
+		WHERE id=$1
+		RETURNING completed`,
+		task.ID,
+		task.Status,
+		task.Completed,
+		task.StatusTimestampProcessed,
+		task.Timestamp,
+		task.Stdout,
+		task.Stderr)
+	if err != nil {
+		return false, err
+	}
+	if err = tx.Commit(); err != nil {
+		return false, err
+	}
+	return !wasCompleted && updatedCompleted, nil
+}
+
 func handleAgentMessagePostResponse(incoming *map[string]interface{}, uUIDInfo *cachedUUIDInfo) (map[string]interface{}, error) {
 	// got message:
 	/*
@@ -576,7 +614,7 @@ func handleAgentMessagePostResponse(incoming *map[string]interface{}, uUIDInfo *
 		if agentMessage.Responses[i].UserOutput != nil && *agentMessage.Responses[i].UserOutput != "" {
 			// do it in the background - the agent doesn't need the result of this directly
 			//handleAgentMessagePostResponseUserOutput(currentTask, agentResponse, true)
-			enqueueAsyncAgentMessagePostResponse(agentAgentMessagePostResponseChannelMessage{
+			enqueueAsyncAgentMessagePostResponse(agentMessagePostResponseUserOutputChannelMessage{
 				Task:        currentTask,
 				Response:    *agentMessage.Responses[i].UserOutput,
 				SequenceNum: agentMessage.Responses[i].SequenceNumber,
@@ -658,15 +696,11 @@ func handleAgentMessagePostResponse(incoming *map[string]interface{}, uUIDInfo *
 	delete(*incoming, "responses")
 	for _, currentTask := range tasksToUpdate {
 		// always updating at least the timestamp for the last thing that happened
-		_, err = database.DB.NamedExec(`UPDATE task SET
-				status=:status, completed=:completed, status_timestamp_processed=:status_timestamp_processed, "timestamp"=:timestamp,
-				stdout=:stdout, stderr=:stderr
-				WHERE id=:id`, currentTask)
+		transitionedToCompleted, err := updateAgentMessagePostResponseTask(currentTask)
 		if err != nil {
 			logging.LogError(err, "Failed to update task from agent response")
 		}
-		if currentTask.Completed {
-			// use updatedToCompleted to try to make sure we only do this once per task
+		if transitionedToCompleted {
 			go CheckAndProcessTaskCompletionHandlers(currentTask.ID)
 			FileBrowserChannel <- FileBrowserChannelMessage{
 				Task:               &currentTask,
@@ -1933,7 +1967,7 @@ func HandleAgentMessagePostResponseFileBrowser(task databaseStructs.Task, fileBr
 				logging.LogError(err, "failed to marshal filebrowser data to JSON")
 				return
 			}
-			enqueueAsyncAgentMessagePostResponse(agentAgentMessagePostResponseChannelMessage{
+			enqueueAsyncAgentMessagePostResponse(agentMessagePostResponseUserOutputChannelMessage{
 				Task:        task,
 				Response:    string(outputBytes),
 				SequenceNum: nil,
@@ -2859,7 +2893,7 @@ func handleAgentMessagePostResponseCustomBrowser(task databaseStructs.Task, agen
 				logging.LogError(err, "failed to marshal filebrowser data to JSON")
 				return
 			}
-			enqueueAsyncAgentMessagePostResponse(agentAgentMessagePostResponseChannelMessage{
+			enqueueAsyncAgentMessagePostResponse(agentMessagePostResponseUserOutputChannelMessage{
 				Task:        task,
 				Response:    string(outputBytes),
 				SequenceNum: nil,

mythic-docker/src/rabbitmq/utils.go

@@ -767,6 +767,62 @@ func GetCallbackCommandInformation(callback databaseStructs.Callback) []string {
 
 var completionMutex sync.Mutex
 
+func claimTaskCompletionFunctionStart(taskID int, startedColumn string, completedColumn string, functionColumn string, runningStatus PT_TASK_FUNCTION_STATUS) bool {
+	query := fmt.Sprintf(`UPDATE task
+		SET %s=true, status=$2
+		WHERE id=$1
+		AND completed=true
+		AND %s != ''
+		AND %s=false
+		AND %s=false`, startedColumn, functionColumn, startedColumn, completedColumn)
+	result, err := database.DB.Exec(query, taskID, runningStatus)
+	if err != nil {
+		logging.LogError(err, "Failed to claim task completion function start", "task_id", taskID, "started_column", startedColumn)
+		return false
+	}
+	rowsAffected, err := result.RowsAffected()
+	if err != nil {
+		logging.LogError(err, "Failed to check task completion function claim result", "task_id", taskID, "started_column", startedColumn)
+		return false
+	}
+	return rowsAffected == 1
+}
+
+func claimGroupCompletionFunctionStart(parentTaskID int64, groupName string, groupCallbackFunction string) bool {
+	result, err := database.DB.Exec(`UPDATE task
+		SET group_callback_function_started=true, status=$4
+		WHERE parent_task_id=$1
+		AND subtask_group_name=$2
+		AND group_callback_function=$3
+		AND group_callback_function != ''
+		AND completed=true
+		AND group_callback_function_started=false
+		AND group_callback_function_completed=false
+		AND NOT EXISTS (
+			SELECT 1
+			FROM task existing_group_task
+			WHERE existing_group_task.parent_task_id=$1
+			AND existing_group_task.subtask_group_name=$2
+			AND existing_group_task.group_callback_function=$3
+			AND (
+				existing_group_task.completed=false
+				OR existing_group_task.group_callback_function_started=true
+				OR existing_group_task.group_callback_function_completed=true
+			)
+		)`,
+		parentTaskID, groupName, groupCallbackFunction, PT_TASK_FUNCTION_STATUS_GROUP_COMPLETED_FUNCTION)
+	if err != nil {
+		logging.LogError(err, "Failed to claim group completion function start", "parent_task_id", parentTaskID, "group_name", groupName)
+		return false
+	}
+	rowsAffected, err := result.RowsAffected()
+	if err != nil {
+		logging.LogError(err, "Failed to check group completion function claim result", "parent_task_id", parentTaskID, "group_name", groupName)
+		return false
+	}
+	return rowsAffected > 0
+}
+
 // Helper functions for getting information for sending Task data to a container
 func CheckAndProcessTaskCompletionHandlers(taskId int) {
 	// check if this task has a completion function
@@ -783,8 +839,11 @@ func CheckAndProcessTaskCompletionHandlers(taskId int) {
 	err := database.DB.Get(&task, `SELECT
 		task.parent_task_id, task.operator_id, task.completed,
 		task.subtask_callback_function, task.subtask_callback_function_completed,
+		task.subtask_callback_function_started,
 		task.group_callback_function, task.group_callback_function_completed, task.completed_callback_function,
-		task.completed_callback_function_completed, task.subtask_group_name, task.id, task.status, task.eventstepinstance_id
+		task.group_callback_function_started,
+		task.completed_callback_function_completed, task.completed_callback_function_started,
+		task.subtask_group_name, task.id, task.status, task.eventstepinstance_id
 		FROM task
 		WHERE task.id=$1`, taskId)
 	if err != nil {
@@ -799,18 +858,23 @@ func CheckAndProcessTaskCompletionHandlers(taskId int) {
 	}
 	if task.ParentTaskID.Valid {
 		err = database.DB.Get(&parentTask, `SELECT 
-    		task.id, task.status, task.completed, task.eventstepinstance_id, task.completed_callback_function_completed,
-    		task.subtask_callback_function,
-    		c.script_only "command.script_only"
-    		from task
-    		LEFT OUTER JOIN command c on task.command_id = c.id
-    		WHERE task.id=$1`, task.ParentTaskID.Int64)
+		task.id, task.status, task.completed, task.eventstepinstance_id, task.completed_callback_function, task.completed_callback_function_completed,
+		task.completed_callback_function_started,
+		task.subtask_callback_function,
+		c.script_only "command.script_only"
+		from task
+		LEFT OUTER JOIN command c on task.command_id = c.id
+		WHERE task.id=$1`, task.ParentTaskID.Int64)
 		if err != nil {
 			logging.LogError(err, "Failed to get parent task information")
 		}
 	}
 	// now we have info for both the current task that finished something and the parent task (if there is one)
 	if task.CompletedCallbackFunction != "" && !task.CompletedCallbackFunctionCompleted {
+		if !claimTaskCompletionFunctionStart(task.ID, "completed_callback_function_started", "completed_callback_function_completed",
+			"completed_callback_function", PT_TASK_FUNCTION_STATUS_COMPLETION_FUNCTION) {
+			return
+		}
 		// this task has a completion function set, and it hasn't been executed, so run it
 		// ex in create_tasking: completionFunctionName := "shellCompleted"
 		//	response.CompletionFunctionName = &completionFunctionName
@@ -820,9 +884,7 @@ func CheckAndProcessTaskCompletionHandlers(taskId int) {
 			CompletionFunctionName: task.CompletedCallbackFunction,
 		}
 		task.Status = PT_TASK_FUNCTION_STATUS_COMPLETION_FUNCTION
-		if _, err := database.DB.NamedExec(`UPDATE task SET status=:status WHERE id=:id`, task); err != nil {
-			logging.LogError(err, "Failed to update status to completion task running")
-		} else if err = RabbitMQConnection.SendPtTaskCompletionFunction(taskMessage); err != nil {
+		if err = RabbitMQConnection.SendPtTaskCompletionFunction(taskMessage); err != nil {
 			logging.LogError(err, "Failed to send task completion function message to container")
 			task.Status = PT_TASK_FUNCTION_STATUS_COMPLETION_FUNCTION_ERROR
 			if _, err = database.DB.NamedExec(`UPDATE task SET status=:status WHERE id=:id`, task); err != nil {
@@ -830,6 +892,14 @@ func CheckAndProcessTaskCompletionHandlers(taskId int) {
 			}
 		}
 	} else if task.SubtaskCallbackFunction != "" && !task.SubtaskCallbackFunctionCompleted {
+		if !task.ParentTaskID.Valid {
+			logging.LogError(nil, "Task has subtask callback function but no parent task", "task_id", task.ID)
+			return
+		}
+		if !claimTaskCompletionFunctionStart(task.ID, "subtask_callback_function_started", "subtask_callback_function_completed",
+			"subtask_callback_function", PT_TASK_FUNCTION_STATUS_SUBTASK_COMPLETED_FUNCTION) {
+			return
+		}
 		// this task has a subtask callback function, and that subtask function hasn't completed successfully
 		// ex in create_tasking:
 		//	completionFunctionName := "pwdCompleted"
@@ -854,6 +924,10 @@ func CheckAndProcessTaskCompletionHandlers(taskId int) {
 			}
 		}
 	} else if task.SubtaskGroupName != "" && task.GroupCallbackFunction != "" && !task.GroupCallbackFunctionCompleted {
+		if !task.ParentTaskID.Valid {
+			logging.LogError(nil, "Task has group callback function but no parent task", "task_id", task.ID)
+			return
+		}
 		// we have a subtask group name, we have a group callback function defined, and that group callback function is done
 		// need to check if we're the last one in the group to finish - if so, we need to call the function, if not do nothing
 		// this function is executed for the PARENT_TASK
@@ -872,6 +946,9 @@ func CheckAndProcessTaskCompletionHandlers(taskId int) {
 			task.ParentTaskID.Int64, task.SubtaskGroupName); err != nil {
 			logging.LogError(err, "Failed to fetch other group tasks")
 		} else if len(groupTasks) == 0 {
+			if !claimGroupCompletionFunctionStart(task.ParentTaskID.Int64, task.SubtaskGroupName, task.GroupCallbackFunction) {
+				return
+			}
 			logging.LogInfo("subtask completed, parent task completion function going to run", "taskId", task.ParentTaskID.Int64, "subtaskid", task.ID)
 			err = RabbitMQConnection.SendPtTaskCompletionFunction(taskMessage)
 			if err != nil {
@@ -901,6 +978,10 @@ func CheckAndProcessTaskCompletionHandlers(taskId int) {
 						logging.LogError(err, "Failed to update parent task information to submitted")
 					}
 					if parentTask.CompletedCallbackFunction != "" && !parentTask.CompletedCallbackFunctionCompleted {
+						if !claimTaskCompletionFunctionStart(parentTask.ID, "completed_callback_function_started", "completed_callback_function_completed",
+							"completed_callback_function", PT_TASK_FUNCTION_STATUS_COMPLETION_FUNCTION) {
+							return
+						}
 						// this task has a completion function set, and it hasn't been executed, so run it
 						// ex in create_tasking: completionFunctionName := "shellCompleted"
 						//	response.CompletionFunctionName = &completionFunctionName
@@ -912,8 +993,8 @@ func CheckAndProcessTaskCompletionHandlers(taskId int) {
 						err = RabbitMQConnection.SendPtTaskCompletionFunction(taskMessage)
 						if err != nil {
 							logging.LogError(err, "Failed to send task completion function message to container")
-							task.Status = PT_TASK_FUNCTION_STATUS_COMPLETION_FUNCTION_ERROR
-							if _, err = database.DB.NamedExec(`UPDATE task SET status=:status WHERE id=:id`, task); err != nil {
+							parentTask.Status = PT_TASK_FUNCTION_STATUS_COMPLETION_FUNCTION_ERROR
+							if _, err = database.DB.NamedExec(`UPDATE task SET status=:status WHERE id=:id`, parentTask); err != nil {
 								logging.LogError(err, "Failed to update task status for completion function error")
 							}
 						}

mythic-docker/src/rabbitmq/utils_proxy_traffic.go

@@ -13,6 +13,7 @@ import (
 	"sort"
 	"strconv"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/its-a-feature/Mythic/database"
@@ -28,6 +29,7 @@ const (
 	CALLBACK_PORT_TYPE_SOCKS       CallbackPortType = "socks"
 	CALLBACK_PORT_TYPE_RPORTFWD                     = "rpfwd"
 	CALLBACK_PORT_TYPE_INTERACTIVE                  = "interactive"
+	callbackPortByteFlushInterval                   = 20 * time.Second
 )
 
 type proxyToAgentMessage struct {
@@ -53,37 +55,28 @@ type acceptedConnection struct {
 	TaskUUID                     *string
 	AgentClosedConnection        bool
 }
-type bytesSentToAgentMessage struct {
-	CallbackPortID int `json:"id"`
-	ByteCount      int64
-	Initial        bool
-}
-type bytesReceivedFromAgentMessage struct {
-	CallbackPortID int `json:"id"`
-	ByteCount      int64
-	Initial        bool
-}
-
 type callbackPortUsage struct {
-	CallbackPortID             int              `json:"id" db:"id"`
-	CallbackID                 int              `json:"callback_id" db:"callback_id"`
-	CallbackDisplayID          int              `json:"callback_display_id" db:"callback_display_id"`
-	TaskID                     int              `json:"task_id" db:"task_id"`
-	LocalPort                  int              `json:"local_port" db:"local_port"`
-	RemotePort                 int              `json:"remote_port" db:"remote_port"`
-	RemoteIP                   string           `json:"remote_ip" db:"remote_ip"`
-	OperationID                int              `json:"operation_id" db:"operation_id"`
-	Username                   string           `json:"username" db:"username"`
-	Password                   string           `json:"password" db:"password"`
-	PortType                   CallbackPortType `json:"port_type" db:"port_type"`
-	listener                   *net.Listener
-	bytesReceivedFromAgentChan chan bytesReceivedFromAgentMessage
-	bytesSentToAgentChan       chan bytesSentToAgentMessage
-	acceptedConnections        *[]*acceptedConnection
-	messagesToAgent            chan proxyToAgentMessage
-	interactiveMessagesToAgent chan agentMessagePostResponseInteractive
-	newConnectionChannel       chan *acceptedConnection
-	removeConnectionsChannel   chan *acceptedConnection
+	CallbackPortID                    int              `json:"id" db:"id"`
+	CallbackID                        int              `json:"callback_id" db:"callback_id"`
+	CallbackDisplayID                 int              `json:"callback_display_id" db:"callback_display_id"`
+	TaskID                            int              `json:"task_id" db:"task_id"`
+	LocalPort                         int              `json:"local_port" db:"local_port"`
+	RemotePort                        int              `json:"remote_port" db:"remote_port"`
+	RemoteIP                          string           `json:"remote_ip" db:"remote_ip"`
+	OperationID                       int              `json:"operation_id" db:"operation_id"`
+	Username                          string           `json:"username" db:"username"`
+	Password                          string           `json:"password" db:"password"`
+	PortType                          CallbackPortType `json:"port_type" db:"port_type"`
+	listener                          *net.Listener
+	bytesReceivedFromAgent            atomic.Int64
+	bytesSentToAgent                  atomic.Int64
+	lastFlushedBytesReceivedFromAgent atomic.Int64
+	lastFlushedBytesSentToAgent       atomic.Int64
+	acceptedConnections               *[]*acceptedConnection
+	messagesToAgent                   chan proxyToAgentMessage
+	interactiveMessagesToAgent        chan agentMessagePostResponseInteractive
+	newConnectionChannel              chan *acceptedConnection
+	removeConnectionsChannel          chan *acceptedConnection
 	// messagesFromAgent - these get parsed by manageConnections and passed to the right connection's messagesFromAgent
 	messagesFromAgent            chan proxyFromAgentMessage
 	interactiveMessagesFromAgent chan agentMessagePostResponseInteractive
@@ -96,8 +89,6 @@ type callbackPortsInUse struct {
 	portsByCallbackIDAndType     map[int]map[CallbackPortType][]*callbackPortUsage
 	callbacksWithPorts           []int
 	proxyFromAgentMessageChannel chan ProxyFromAgentMessageForMythic
-	bytesReceivedFromAgentChan   chan bytesReceivedFromAgentMessage
-	bytesSentToAgentChan         chan bytesSentToAgentMessage
 	sync.RWMutex
 }
 
@@ -227,8 +218,6 @@ func (c *callbackPortsInUse) Initialize() {
 	c.ports = make([]*callbackPortUsage, 0)
 	c.resetPortIndexesLocked()
 	c.proxyFromAgentMessageChannel = make(chan ProxyFromAgentMessageForMythic, 2000)
-	c.bytesReceivedFromAgentChan = make(chan bytesReceivedFromAgentMessage, 2000)
-	c.bytesSentToAgentChan = make(chan bytesSentToAgentMessage, 2000)
 	c.Unlock()
 	go c.ListenForProxyFromAgentMessage()
 	go c.ListenForNewByteTransferUpdates()
@@ -236,7 +225,7 @@ func (c *callbackPortsInUse) Initialize() {
     	callbackport.id, callbackport.callback_id, callbackport.task_id, 
     	callbackport.local_port, callbackport.remote_port, callbackport.remote_ip,
     	callbackport.port_type, callbackport.operation_id, callbackport.username,
-    	callbackport.password,
+        callbackport.password, callbackport.bytes_sent, callbackport.bytes_received,
     	callback.display_id "callback.display_id"
 		FROM callbackport 
 		JOIN callback ON callbackport.callback_id=callback.id
@@ -244,16 +233,6 @@ func (c *callbackPortsInUse) Initialize() {
 		logging.LogError(err, "Failed to load callback ports from database")
 	} else {
 		for _, proxy := range callbackPorts {
-			c.bytesSentToAgentChan <- bytesSentToAgentMessage{
-				ByteCount:      proxy.BytesSent,
-				CallbackPortID: proxy.ID,
-				Initial:        true,
-			}
-			c.bytesReceivedFromAgentChan <- bytesReceivedFromAgentMessage{
-				ByteCount:      proxy.BytesReceived,
-				CallbackPortID: proxy.ID,
-				Initial:        true,
-			}
 			newPort := callbackPortUsage{
 				CallbackPortID:               proxy.ID,
 				CallbackID:                   proxy.CallbackID,
@@ -266,8 +245,6 @@ func (c *callbackPortsInUse) Initialize() {
 				OperationID:                  proxy.OperationID,
 				Username:                     proxy.Username,
 				Password:                     proxy.Password,
-				bytesSentToAgentChan:         c.bytesSentToAgentChan,
-				bytesReceivedFromAgentChan:   c.bytesReceivedFromAgentChan,
 				messagesToAgent:              make(chan proxyToAgentMessage, 1000),
 				newConnectionChannel:         make(chan *acceptedConnection, 1000),
 				removeConnectionsChannel:     make(chan *acceptedConnection, 1000),
@@ -276,6 +253,7 @@ func (c *callbackPortsInUse) Initialize() {
 				interactiveMessagesFromAgent: make(chan agentMessagePostResponseInteractive, 1000),
 				stopAllConnections:           make(chan bool),
 			}
+			newPort.initializeByteCounters(proxy.BytesSent, proxy.BytesReceived)
 			acceptedConnections := make([]*acceptedConnection, 0)
 			newPort.acceptedConnections = &acceptedConnections
 			if err := newPort.Start(); err != nil {
@@ -374,6 +352,107 @@ func cloneCallbackPortUsages(ports []*callbackPortUsage) []*callbackPortUsage {
 	return append([]*callbackPortUsage(nil), ports...)
 }
 
+func clampCallbackPortByteCount(value int64) int64 {
+	if value < 0 {
+		return 0
+	}
+	if value >= POSTGRES_MAX_BIGINT {
+		return POSTGRES_MAX_BIGINT - 1
+	}
+	return value
+}
+
+func addCallbackPortByteCounts(value int64, delta int64) int64 {
+	value = clampCallbackPortByteCount(value)
+	if delta <= 0 {
+		return value
+	}
+	if delta >= POSTGRES_MAX_BIGINT-value {
+		return POSTGRES_MAX_BIGINT - 1
+	}
+	return value + delta
+}
+
+func addClampedAtomicInt64(counter *atomic.Int64, delta int64) {
+	if delta <= 0 {
+		return
+	}
+	for {
+		current := counter.Load()
+		next := addCallbackPortByteCounts(current, delta)
+		if next == current || counter.CompareAndSwap(current, next) {
+			return
+		}
+	}
+}
+
+func (p *callbackPortUsage) initializeByteCounters(bytesSentToAgent int64, bytesReceivedFromAgent int64) {
+	bytesSentToAgent = clampCallbackPortByteCount(bytesSentToAgent)
+	bytesReceivedFromAgent = clampCallbackPortByteCount(bytesReceivedFromAgent)
+	p.bytesSentToAgent.Store(bytesSentToAgent)
+	p.lastFlushedBytesSentToAgent.Store(bytesSentToAgent)
+	p.bytesReceivedFromAgent.Store(bytesReceivedFromAgent)
+	p.lastFlushedBytesReceivedFromAgent.Store(bytesReceivedFromAgent)
+}
+
+func (p *callbackPortUsage) applyPersistedByteCounters(bytesSentToAgent int64, bytesReceivedFromAgent int64) {
+	applyPersistedByteCounter(&p.bytesSentToAgent, &p.lastFlushedBytesSentToAgent, bytesSentToAgent)
+	applyPersistedByteCounter(&p.bytesReceivedFromAgent, &p.lastFlushedBytesReceivedFromAgent, bytesReceivedFromAgent)
+}
+
+func applyPersistedByteCounter(counter *atomic.Int64, lastFlushed *atomic.Int64, persistedValue int64) {
+	persistedValue = clampCallbackPortByteCount(persistedValue)
+	for {
+		currentDelta := counter.Load()
+		next := addCallbackPortByteCounts(persistedValue, currentDelta)
+		if counter.CompareAndSwap(currentDelta, next) {
+			lastFlushed.Store(persistedValue)
+			return
+		}
+	}
+}
+
+func (p *callbackPortUsage) addBytesSentToAgent(byteCount int64) {
+	addClampedAtomicInt64(&p.bytesSentToAgent, byteCount)
+}
+
+func (p *callbackPortUsage) addBytesReceivedFromAgent(byteCount int64) {
+	addClampedAtomicInt64(&p.bytesReceivedFromAgent, byteCount)
+}
+
+func (p *callbackPortUsage) flushByteCounters() {
+	if p.CallbackPortID <= 0 {
+		return
+	}
+	currentBytesSentToAgent := p.bytesSentToAgent.Load()
+	bytesSentToAgent := clampCallbackPortByteCount(currentBytesSentToAgent)
+	if bytesSentToAgent != currentBytesSentToAgent {
+		p.bytesSentToAgent.CompareAndSwap(currentBytesSentToAgent, bytesSentToAgent)
+	}
+	if bytesSentToAgent != p.lastFlushedBytesSentToAgent.Load() &&
+		updateCallbackPortStats("bytes_sent", bytesSentToAgent, p.CallbackPortID) {
+		p.lastFlushedBytesSentToAgent.Store(bytesSentToAgent)
+	}
+	currentBytesReceivedFromAgent := p.bytesReceivedFromAgent.Load()
+	bytesReceivedFromAgent := clampCallbackPortByteCount(currentBytesReceivedFromAgent)
+	if bytesReceivedFromAgent != currentBytesReceivedFromAgent {
+		p.bytesReceivedFromAgent.CompareAndSwap(currentBytesReceivedFromAgent, bytesReceivedFromAgent)
+	}
+	if bytesReceivedFromAgent != p.lastFlushedBytesReceivedFromAgent.Load() &&
+		updateCallbackPortStats("bytes_received", bytesReceivedFromAgent, p.CallbackPortID) {
+		p.lastFlushedBytesReceivedFromAgent.Store(bytesReceivedFromAgent)
+	}
+}
+
+func (c *callbackPortsInUse) flushChangedByteCounters() {
+	c.RLock()
+	ports := cloneCallbackPortUsages(c.ports)
+	c.RUnlock()
+	for _, port := range ports {
+		port.flushByteCounters()
+	}
+}
+
 func (c *callbackPortsInUse) getPortsForCallbackAndType(callbackID int, portType CallbackPortType) []*callbackPortUsage {
 	c.RLock()
 	defer c.RUnlock()
@@ -431,83 +510,28 @@ func (c *callbackPortsInUse) findPortForRemoval(callbackId int, portType Callbac
 	return nil
 }
 func (c *callbackPortsInUse) ListenForNewByteTransferUpdates() {
-	oldByteValues := make(map[int]map[string]int64)
-	currentByteValues := make(map[int]map[string]int64)
-	updateValuesChan := make(chan bool)
-	go func() {
-		// every 30s send a signal to make sure we update our current byte values if they've changed
-		for {
-			select {
-			case <-time.After(20 * time.Second):
-				updateValuesChan <- true
-			}
-		}
-	}()
+	ticker := time.NewTicker(callbackPortByteFlushInterval)
+	defer ticker.Stop()
 	defer func() {
 		logging.LogError(nil, "no longer listening for new byte transfer updates")
 	}()
-	for {
-		select {
-		case <-updateValuesChan:
-			for callbackPortID, _ := range currentByteValues {
-				if currentByteValues[callbackPortID]["sent"] != oldByteValues[callbackPortID]["sent"] {
-					go updateCallbackPortStats("bytes_sent", currentByteValues[callbackPortID]["sent"], callbackPortID)
-					oldByteValues[callbackPortID]["sent"] = currentByteValues[callbackPortID]["sent"]
-				}
-				if currentByteValues[callbackPortID]["received"] != oldByteValues[callbackPortID]["received"] {
-					go updateCallbackPortStats("bytes_received", currentByteValues[callbackPortID]["received"], callbackPortID)
-					oldByteValues[callbackPortID]["received"] = currentByteValues[callbackPortID]["received"]
-				}
-			}
-
-		case bytesFromAgentMsg := <-c.bytesReceivedFromAgentChan:
-			if _, ok := currentByteValues[bytesFromAgentMsg.CallbackPortID]; !ok {
-				currentByteValues[bytesFromAgentMsg.CallbackPortID] = map[string]int64{
-					"received": 0,
-					"sent":     0,
-				}
-				oldByteValues[bytesFromAgentMsg.CallbackPortID] = map[string]int64{
-					"received": 0,
-					"sent":     0,
-				}
-			}
-			if bytesFromAgentMsg.Initial {
-				oldByteValues[bytesFromAgentMsg.CallbackPortID]["received"] = bytesFromAgentMsg.ByteCount
-				currentByteValues[bytesFromAgentMsg.CallbackPortID]["received"] = bytesFromAgentMsg.ByteCount
-			} else {
-				currentByteValues[bytesFromAgentMsg.CallbackPortID]["received"] += bytesFromAgentMsg.ByteCount
-			}
-		case bytesSentToAgentMsg := <-c.bytesSentToAgentChan:
-			if _, ok := currentByteValues[bytesSentToAgentMsg.CallbackPortID]; !ok {
-				currentByteValues[bytesSentToAgentMsg.CallbackPortID] = map[string]int64{
-					"received": 0,
-					"sent":     0,
-				}
-				oldByteValues[bytesSentToAgentMsg.CallbackPortID] = map[string]int64{
-					"received": 0,
-					"sent":     0,
-				}
-			}
-			if bytesSentToAgentMsg.Initial {
-				oldByteValues[bytesSentToAgentMsg.CallbackPortID]["sent"] = bytesSentToAgentMsg.ByteCount
-				currentByteValues[bytesSentToAgentMsg.CallbackPortID]["sent"] = bytesSentToAgentMsg.ByteCount
-			} else {
-				currentByteValues[bytesSentToAgentMsg.CallbackPortID]["sent"] += bytesSentToAgentMsg.ByteCount
-			}
-		}
+	for range ticker.C {
+		c.flushChangedByteCounters()
 	}
 }
 
-func updateCallbackPortStats(field string, value int64, callbackPortID int) {
-	updatedValue := value
-	if updatedValue > POSTGRES_MAX_BIGINT {
-		updatedValue = POSTGRES_MAX_BIGINT - 1
+func updateCallbackPortStats(field string, value int64, callbackPortID int) bool {
+	if callbackPortID <= 0 {
+		return false
 	}
+	updatedValue := clampCallbackPortByteCount(value)
 	_, err := database.DB.Exec(fmt.Sprintf("UPDATE callbackport SET %s=$1 WHERE id=$2",
-		field), value, callbackPortID)
+		field), updatedValue, callbackPortID)
 	if err != nil {
 		logging.LogError(err, "Failed to update callback port stats")
+		return false
 	}
+	return true
 }
 func (c *callbackPortsInUse) ListenForProxyFromAgentMessage() {
 	for {
@@ -712,18 +736,6 @@ func (c *callbackPortsInUse) GetDataForCallbackId(callbackId int, portType strin
 func (c *callbackPortsInUse) Add(callbackId int, portType CallbackPortType, localPort int, remotePort int,
 	remoteIP string, taskId int, operationId int, bytesSentToAgent int64, bytesReceivedFromAgent int64,
 	callbackPortID int, username string, password string) error {
-	if callbackPortID > 0 {
-		c.bytesSentToAgentChan <- bytesSentToAgentMessage{
-			ByteCount:      bytesSentToAgent,
-			CallbackPortID: callbackPortID,
-			Initial:        true,
-		}
-		c.bytesReceivedFromAgentChan <- bytesReceivedFromAgentMessage{
-			ByteCount:      bytesReceivedFromAgent,
-			CallbackPortID: callbackPortID,
-			Initial:        true,
-		}
-	}
 	newPort := callbackPortUsage{
 		CallbackPortID:               callbackPortID,
 		CallbackID:                   callbackId,
@@ -735,8 +747,6 @@ func (c *callbackPortsInUse) Add(callbackId int, portType CallbackPortType, loca
 		PortType:                     portType,
 		Username:                     username,
 		Password:                     password,
-		bytesSentToAgentChan:         c.bytesSentToAgentChan,
-		bytesReceivedFromAgentChan:   c.bytesReceivedFromAgentChan,
 		messagesToAgent:              make(chan proxyToAgentMessage, 1000),
 		newConnectionChannel:         make(chan *acceptedConnection, 1000),
 		removeConnectionsChannel:     make(chan *acceptedConnection, 1000),
@@ -745,6 +755,11 @@ func (c *callbackPortsInUse) Add(callbackId int, portType CallbackPortType, loca
 		interactiveMessagesFromAgent: make(chan agentMessagePostResponseInteractive, 1000),
 		stopAllConnections:           make(chan bool),
 	}
+	if callbackPortID > 0 {
+		newPort.initializeByteCounters(bytesSentToAgent, bytesReceivedFromAgent)
+	} else {
+		newPort.initializeByteCounters(0, 0)
+	}
 	acceptedConnections := make([]*acceptedConnection, 0)
 	newPort.acceptedConnections = &acceptedConnections
 	callbackPort := databaseStructs.Callbackport{}
@@ -759,7 +774,7 @@ func (c *callbackPortsInUse) Add(callbackId int, portType CallbackPortType, loca
 		logging.LogError(err, "Failed to start new proxy port")
 		return err
 	}
-	err = database.DB.Get(&callbackPort, `SELECT id FROM callbackport WHERE
+	err = database.DB.Get(&callbackPort, `SELECT id, bytes_sent, bytes_received FROM callbackport WHERE
                                 operation_id=$1 AND callback_id=$2 AND local_port=$3 AND port_type=$4
                                 AND remote_ip=$5 AND remote_port=$6 AND username=$7 AND password=$8`,
 		operationId, callbackId, localPort, portType, remoteIP, remotePort, username, password)
@@ -775,7 +790,7 @@ func (c *callbackPortsInUse) Add(callbackId int, portType CallbackPortType, loca
 			}
 			return err
 		}
-		err = statement.Get(&newPort.CallbackPortID, newPort)
+		err = statement.Get(&newPort.CallbackPortID, &newPort)
 		if err != nil {
 			logging.LogError(err, "Failed to insert new callback port")
 			if err := newPort.Stop(); err != nil {
@@ -786,6 +801,9 @@ func (c *callbackPortsInUse) Add(callbackId int, portType CallbackPortType, loca
 	} else if err == nil {
 		_, err = database.DB.NamedExec(`UPDATE callbackport SET deleted=false WHERE id=:id`, callbackPort)
 		newPort.CallbackPortID = callbackPort.ID
+		if callbackPortID <= 0 {
+			newPort.applyPersistedByteCounters(callbackPort.BytesSent, callbackPort.BytesReceived)
+		}
 	} else {
 		logging.LogError(err, "Failed to create new callback port mapping")
 		if err := newPort.Stop(); err != nil {
@@ -829,6 +847,7 @@ func (c *callbackPortsInUse) Remove(callbackId int, portType CallbackPortType, l
 		logging.LogError(err, "Failed to stop proxy")
 		return err
 	}
+	port.flushByteCounters()
 	queryString := `UPDATE callbackport SET deleted=true WHERE id=$1`
 	queryArgs := []interface{}{port.CallbackPortID}
 	_, err = database.DB.Exec(queryString, queryArgs...)
@@ -1412,12 +1431,7 @@ func (p *callbackPortUsage) readFromAgentSocksConn(newConnection *acceptedConnec
 				p.removeConnectionsChannel <- newConnection
 				return
 			}
-			// non-blocking send stats update
-			go func(byteCount int64, callbackPortID int) {
-				select {
-				case p.bytesReceivedFromAgentChan <- bytesReceivedFromAgentMessage{ByteCount: byteCount, CallbackPortID: callbackPortID, Initial: false}:
-				}
-			}(int64(len(dataBytes)), p.CallbackPortID)
+			p.addBytesReceivedFromAgent(int64(len(dataBytes)))
 
 			if agentMessage.IsExit {
 				logging.LogDebug("got message from agent isExit", "server_id", newConnection.ServerID)
@@ -1746,12 +1760,7 @@ func (p *callbackPortUsage) handleSocksConnections() {
 											CallbackID:      p.CallbackID,
 											ProxyType:       p.PortType,
 										}
-										// non-blocking send stats update
-										go func(byteCount int64, callbackPortID int) {
-											select {
-											case p.bytesSentToAgentChan <- bytesSentToAgentMessage{ByteCount: byteCount, CallbackPortID: callbackPortID, Initial: false}:
-											}
-										}(int64(newUDPConnectionBufLength), p.CallbackPortID)
+										p.addBytesSentToAgent(int64(newUDPConnectionBufLength))
 										go func(remoteAddr net.Addr) {
 											// work on this section!!
 											// handle the read from agent and write back out to socket here so we can use the same addr
@@ -1774,12 +1783,7 @@ func (p *callbackPortUsage) handleSocksConnections() {
 														p.removeConnectionsChannel <- &newUDPConnection
 														return
 													}
-													// non-blocking send stats update
-													go func(byteCount int64, callbackPortID int) {
-														select {
-														case p.bytesReceivedFromAgentChan <- bytesReceivedFromAgentMessage{ByteCount: byteCount, CallbackPortID: callbackPortID, Initial: false}:
-														}
-													}(int64(len(dataBytes)), p.CallbackPortID)
+													p.addBytesReceivedFromAgent(int64(len(dataBytes)))
 
 													if agentMessage.IsExit {
 														//logging.LogDebug("got message from agent isExit", "server_id", newConnection.ServerID)
@@ -1821,12 +1825,7 @@ func (p *callbackPortUsage) handleSocksConnections() {
 							CallbackID:      p.CallbackID,
 							ProxyType:       p.PortType,
 						}
-						// non-blocking send stats update
-						go func(byteCount int64, callbackPortID int) {
-							select {
-							case p.bytesSentToAgentChan <- bytesSentToAgentMessage{ByteCount: byteCount, CallbackPortID: callbackPortID, Initial: false}:
-							}
-						}(int64(length), p.CallbackPortID)
+						p.addBytesSentToAgent(int64(length))
 					}
 					if err != nil {
 						logging.LogError(err, "closing tcp connection")
@@ -1866,11 +1865,7 @@ func (p *callbackPortUsage) handleRpfwdConnections(newConnection *acceptedConnec
 					p.removeConnectionsChannel <- newConnection
 					return
 				}
-				// non-blocking send stats update
-				select {
-				case p.bytesReceivedFromAgentChan <- bytesReceivedFromAgentMessage{ByteCount: int64(len(dataBytes)), CallbackPortID: p.CallbackPortID}:
-				default:
-				}
+				p.addBytesReceivedFromAgent(int64(len(dataBytes)))
 				if agentMessage.IsExit {
 					//logging.LogDebug("got message isExit", "server_id", newConnection.ServerID)
 					// cleanup the connection data, but don't send an exit back to the agent
@@ -1906,11 +1901,7 @@ func (p *callbackPortUsage) handleRpfwdConnections(newConnection *acceptedConnec
 					ProxyType:       p.PortType,
 					CallbackID:      p.CallbackID,
 				}
-				// non-blocking send stats update
-				select {
-				case p.bytesSentToAgentChan <- bytesSentToAgentMessage{ByteCount: int64(length), CallbackPortID: p.CallbackPortID}:
-				default:
-				}
+				p.addBytesSentToAgent(int64(length))
 				//fmt.Printf("Message sent to p.messagesToAgent channel for chan %d\n", newConnection.ServerID)
 			}
 			if err != nil {
@@ -1978,11 +1969,7 @@ func (p *callbackPortUsage) handleInteractiveConnections() {
 								p.removeConnectionsChannel <- &newConnection
 								return
 							}
-							// non-blocking send stats update
-							select {
-							case p.bytesReceivedFromAgentChan <- bytesReceivedFromAgentMessage{ByteCount: int64(len(dataBytes)), CallbackPortID: p.CallbackPortID}:
-							default:
-							}
+							p.addBytesReceivedFromAgent(int64(len(dataBytes)))
 							if agentMessage.MessageType == InteractiveTask.Exit {
 								//logging.LogDebug("got message isExit", "server_id", newConnection.ServerID)
 								// cleanup the connection data
@@ -2018,11 +2005,7 @@ func (p *callbackPortUsage) handleInteractiveConnections() {
 								CallbackID:                 p.CallbackID,
 								ProxyType:                  p.PortType,
 							}
-							// non-blocking send stats update
-							select {
-							case p.bytesSentToAgentChan <- bytesSentToAgentMessage{ByteCount: int64(length), CallbackPortID: p.CallbackPortID}:
-							default:
-							}
+							p.addBytesSentToAgent(int64(length))
 							//fmt.Printf("Message sent to p.messagesToAgent channel for chan %d\n", newConnection.ServerID)
 						}
 

mythic-docker/src/rabbitmq/utils_proxy_traffic_test.go

@@ -2,6 +2,7 @@ package rabbitmq
 
 import (
 	"slices"
+	"sync"
 	"testing"
 )
 
@@ -92,3 +93,90 @@ func TestCallbackPortsInUseCombinedDataFetchUsesIndexedPorts(t *testing.T) {
 		t.Fatalf("expected proxy data to be drained, got %#v", proxyData)
 	}
 }
+
+func TestCallbackPortUsageByteCountersAccumulateWithoutChannels(t *testing.T) {
+	port := newTestCallbackPort(10, CALLBACK_PORT_TYPE_SOCKS, 7001, 1)
+	port.initializeByteCounters(10, 20)
+
+	port.addBytesSentToAgent(5)
+	port.addBytesReceivedFromAgent(7)
+
+	if got := port.bytesSentToAgent.Load(); got != 15 {
+		t.Fatalf("expected bytes sent to agent to be 15, got %d", got)
+	}
+	if got := port.bytesReceivedFromAgent.Load(); got != 27 {
+		t.Fatalf("expected bytes received from agent to be 27, got %d", got)
+	}
+	if got := port.lastFlushedBytesSentToAgent.Load(); got != 10 {
+		t.Fatalf("expected last flushed bytes sent to agent to remain 10, got %d", got)
+	}
+	if got := port.lastFlushedBytesReceivedFromAgent.Load(); got != 20 {
+		t.Fatalf("expected last flushed bytes received from agent to remain 20, got %d", got)
+	}
+}
+
+func TestCallbackPortUsageByteCountersApplyPersistedBaseline(t *testing.T) {
+	port := newTestCallbackPort(10, CALLBACK_PORT_TYPE_SOCKS, 7001, 1)
+	port.initializeByteCounters(0, 0)
+	port.addBytesSentToAgent(3)
+	port.addBytesReceivedFromAgent(4)
+
+	port.applyPersistedByteCounters(100, 200)
+
+	if got := port.bytesSentToAgent.Load(); got != 103 {
+		t.Fatalf("expected bytes sent to agent to include persisted baseline and pending delta, got %d", got)
+	}
+	if got := port.bytesReceivedFromAgent.Load(); got != 204 {
+		t.Fatalf("expected bytes received from agent to include persisted baseline and pending delta, got %d", got)
+	}
+	if got := port.lastFlushedBytesSentToAgent.Load(); got != 100 {
+		t.Fatalf("expected last flushed bytes sent to agent to track persisted baseline, got %d", got)
+	}
+	if got := port.lastFlushedBytesReceivedFromAgent.Load(); got != 200 {
+		t.Fatalf("expected last flushed bytes received from agent to track persisted baseline, got %d", got)
+	}
+}
+
+func TestCallbackPortUsageByteCountersClampAtPostgresBigint(t *testing.T) {
+	port := newTestCallbackPort(10, CALLBACK_PORT_TYPE_SOCKS, 7001, 1)
+	port.initializeByteCounters(POSTGRES_MAX_BIGINT-2, POSTGRES_MAX_BIGINT-3)
+
+	port.addBytesSentToAgent(10)
+	port.addBytesReceivedFromAgent(10)
+
+	if got := port.bytesSentToAgent.Load(); got != POSTGRES_MAX_BIGINT-1 {
+		t.Fatalf("expected bytes sent to agent to clamp to %d, got %d", POSTGRES_MAX_BIGINT-1, got)
+	}
+	if got := port.bytesReceivedFromAgent.Load(); got != POSTGRES_MAX_BIGINT-1 {
+		t.Fatalf("expected bytes received from agent to clamp to %d, got %d", POSTGRES_MAX_BIGINT-1, got)
+	}
+}
+
+func TestCallbackPortUsageByteCountersHandleConcurrentAdds(t *testing.T) {
+	port := newTestCallbackPort(10, CALLBACK_PORT_TYPE_SOCKS, 7001, 1)
+	port.initializeByteCounters(5, 7)
+
+	const goroutineCount = 32
+	const incrementsPerGoroutine = 1000
+	var wg sync.WaitGroup
+	wg.Add(goroutineCount)
+	for i := 0; i < goroutineCount; i++ {
+		go func() {
+			defer wg.Done()
+			for j := 0; j < incrementsPerGoroutine; j++ {
+				port.addBytesSentToAgent(1)
+				port.addBytesReceivedFromAgent(2)
+			}
+		}()
+	}
+	wg.Wait()
+
+	expectedSent := int64(5 + goroutineCount*incrementsPerGoroutine)
+	expectedReceived := int64(7 + goroutineCount*incrementsPerGoroutine*2)
+	if got := port.bytesSentToAgent.Load(); got != expectedSent {
+		t.Fatalf("expected bytes sent to agent to be %d, got %d", expectedSent, got)
+	}
+	if got := port.bytesReceivedFromAgent.Load(); got != expectedReceived {
+		t.Fatalf("expected bytes received from agent to be %d, got %d", expectedReceived, got)
+	}
+}