admin/trailofbits-buttercup
1
0
Fork 0
mirror of https://github.com/trailofbits/buttercup synced 2026-06-21 14:11:39 +00:00
Code Issues Packages Projects Releases Wiki Activity

Program model integrate (#13)

Browse Source 
* Add documentation

* Add janusgraph custom configuration. Simply graph creation.

* Add properties to edges. Index on uri property.

* reformat

* Create libpng graph in a more reasonable amount of time.

* Update README.md

* Wrap indexing into function

* updating...

* update proto

* fix command

* Update settings for serve and processing indexing

* Add dockerfile to docker container

* Add trigger script

* Fix error with loading graphml file

* reformat

* Remove resource constraints for program model

* Change permissions of temp dir

* Update readme

* Update dockerignore

---------

Co-authored-by: Evan Downing <2077950+evandowning@users.noreply.github.com>
This commit is contained in:
2over12
2025-02-18 11:59:16 -05:00
committed by GitHub
parent 3668cd9301
commit ea2fd093d5
34 changed files with 2706 additions and 366 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.dockerignore
+4 -1
View File
@@ -16,7 +16,10 @@
**/.pdm-build/
**/.pdm-python
**/pdm.lock
**/Dockerfile
fuzzer/dockerfiles/runner_image.Dockerfile
orchestrator/Dockerfile
patcher/Dockerfile
seed-gen/Dockerfile
**/.mypy_cache/
**/.ruff_cache/
common/protos/msg.proto
+18
View File
@@ -90,3 +90,21 @@ message Patch {
string vulnerability_id = 2;
string patch = 3;
}
message IndexRequest {
string package_name = 1;
string sanitizer = 2;
string ossfuzz = 3;
string source_path = 4;
string task_id = 5;
string build_type = 6;
}
message IndexOutput {
string package_name = 1;
string sanitizer = 2;
string ossfuzz = 3;
string source_path = 4;
string task_id = 5;
string build_type = 6;
}
common/src/buttercup/common/datastructures/msg_pb2.py
+5 -1
View File
@@ -13,7 +13,7 @@ _sym_db = _symbol_database.Default()
DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\tmsg.proto\x12\x06msgpb2\"\xec\x01\n\x04Task\x12\x12\n\nmessage_id\x18\x01 \x01(\t\x12\x14\n\x0cmessage_time\x18\x02 \x01(\x03\x12\x0f\n\x07task_id\x18\x03 \x01(\t\x12(\n\ttask_type\x18\x04 \x01(\x0e\x32\x15.msgpb2.Task.TaskType\x12%\n\x07sources\x18\x05 \x03(\x0b\x32\x14.msgpb2.SourceDetail\x12\x10\n\x08\x64\x65\x61\x64line\x18\x06 \x01(\x03\x12\x11\n\tcancelled\x18\x07 \x01(\x08\"3\n\x08TaskType\x12\x12\n\x0eTASK_TYPE_FULL\x10\x00\x12\x13\n\x0fTASK_TYPE_DELTA\x10\x01\"\xb9\x01\n\x0cSourceDetail\x12\x0e\n\x06sha256\x18\x01 \x01(\t\x12\x34\n\x0bsource_type\x18\x02 \x01(\x0e\x32\x1f.msgpb2.SourceDetail.SourceType\x12\x0b\n\x03url\x18\x03 \x01(\t\"V\n\nSourceType\x12\x14\n\x10SOURCE_TYPE_REPO\x10\x00\x12\x1c\n\x18SOURCE_TYPE_FUZZ_TOOLING\x10\x01\x12\x14\n\x10SOURCE_TYPE_DIFF\x10\x02\"*\n\x0cTaskDownload\x12\x1a\n\x04task\x18\x01 \x01(\x0b\x32\x0c.msgpb2.Task\"\'\n\tTaskReady\x12\x1a\n\x04task\x18\x01 \x01(\x0b\x32\x0c.msgpb2.Task\"2\n\nTaskDelete\x12\x0f\n\x07task_id\x18\x01 \x01(\t\x12\x13\n\x0breceived_at\x18\x02 \x01(\x02\"\x92\x01\n\x0c\x42uildRequest\x12\x14\n\x0cpackage_name\x18\x01 \x01(\t\x12\x0e\n\x06\x65ngine\x18\x02 \x01(\t\x12\x11\n\tsanitizer\x18\x03 \x01(\t\x12\x10\n\x08task_dir\x18\x04 \x01(\t\x12\x0f\n\x07task_id\x18\x05 \x01(\t\x12\x12\n\nbuild_type\x18\x06 \x01(\t\x12\x12\n\napply_diff\x18\x07 \x01(\x08\"\x91\x01\n\x0b\x42uildOutput\x12\x14\n\x0cpackage_name\x18\x01 \x01(\t\x12\x0e\n\x06\x65ngine\x18\x02 \x01(\t\x12\x11\n\tsanitizer\x18\x03 \x01(\t\x12\x10\n\x08task_dir\x18\x04 \x01(\t\x12\x0f\n\x07task_id\x18\x05 \x01(\t\x12\x12\n\nbuild_type\x18\x06 \x01(\t\x12\x12\n\napply_diff\x18\x07 \x01(\x08\"^\n\x0fWeightedHarness\x12\x0e\n\x06weight\x18\x01 \x01(\x02\x12\x14\n\x0cpackage_name\x18\x02 \x01(\t\x12\x14\n\x0charness_name\x18\x03 \x01(\t\x12\x0f\n\x07task_id\x18\x04 \x01(\t\"p\n\x05\x43rash\x12#\n\x06target\x18\x01 \x01(\x0b\x32\x13.msgpb2.BuildOutput\x12\x14\n\x0charness_name\x18\x02 \x01(\t\x12\x18\n\x10\x63rash_input_path\x18\x03 \x01(\t\x12\x12\n\nstacktrace\x18\x04 \x01(\t\"G\n\x16\x43onfirmedVulnerability\x12\x1c\n\x05\x63rash\x18\x01 \x01(\x0b\x32\r.msgpb2.Crash\x12\x0f\n\x07vuln_id\x18\x02 \x01(\t\"A\n\x05Patch\x12\x0f\n\x07task_id\x18\x01 \x01(\t\x12\x18\n\x10vulnerability_id\x18\x02 \x01(\t\x12\r\n\x05patch\x18\x03 \x01(\tb\x06proto3')
DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\tmsg.proto\x12\x06msgpb2\"\xec\x01\n\x04Task\x12\x12\n\nmessage_id\x18\x01 \x01(\t\x12\x14\n\x0cmessage_time\x18\x02 \x01(\x03\x12\x0f\n\x07task_id\x18\x03 \x01(\t\x12(\n\ttask_type\x18\x04 \x01(\x0e\x32\x15.msgpb2.Task.TaskType\x12%\n\x07sources\x18\x05 \x03(\x0b\x32\x14.msgpb2.SourceDetail\x12\x10\n\x08\x64\x65\x61\x64line\x18\x06 \x01(\x03\x12\x11\n\tcancelled\x18\x07 \x01(\x08\"3\n\x08TaskType\x12\x12\n\x0eTASK_TYPE_FULL\x10\x00\x12\x13\n\x0fTASK_TYPE_DELTA\x10\x01\"\xb9\x01\n\x0cSourceDetail\x12\x0e\n\x06sha256\x18\x01 \x01(\t\x12\x34\n\x0bsource_type\x18\x02 \x01(\x0e\x32\x1f.msgpb2.SourceDetail.SourceType\x12\x0b\n\x03url\x18\x03 \x01(\t\"V\n\nSourceType\x12\x14\n\x10SOURCE_TYPE_REPO\x10\x00\x12\x1c\n\x18SOURCE_TYPE_FUZZ_TOOLING\x10\x01\x12\x14\n\x10SOURCE_TYPE_DIFF\x10\x02\"*\n\x0cTaskDownload\x12\x1a\n\x04task\x18\x01 \x01(\x0b\x32\x0c.msgpb2.Task\"\'\n\tTaskReady\x12\x1a\n\x04task\x18\x01 \x01(\x0b\x32\x0c.msgpb2.Task\"2\n\nTaskDelete\x12\x0f\n\x07task_id\x18\x01 \x01(\t\x12\x13\n\x0breceived_at\x18\x02 \x01(\x02\"\x92\x01\n\x0c\x42uildRequest\x12\x14\n\x0cpackage_name\x18\x01 \x01(\t\x12\x0e\n\x06\x65ngine\x18\x02 \x01(\t\x12\x11\n\tsanitizer\x18\x03 \x01(\t\x12\x10\n\x08task_dir\x18\x04 \x01(\t\x12\x0f\n\x07task_id\x18\x05 \x01(\t\x12\x12\n\nbuild_type\x18\x06 \x01(\t\x12\x12\n\napply_diff\x18\x07 \x01(\x08\"\x91\x01\n\x0b\x42uildOutput\x12\x14\n\x0cpackage_name\x18\x01 \x01(\t\x12\x0e\n\x06\x65ngine\x18\x02 \x01(\t\x12\x11\n\tsanitizer\x18\x03 \x01(\t\x12\x10\n\x08task_dir\x18\x04 \x01(\t\x12\x0f\n\x07task_id\x18\x05 \x01(\t\x12\x12\n\nbuild_type\x18\x06 \x01(\t\x12\x12\n\napply_diff\x18\x07 \x01(\x08\"^\n\x0fWeightedHarness\x12\x0e\n\x06weight\x18\x01 \x01(\x02\x12\x14\n\x0cpackage_name\x18\x02 \x01(\t\x12\x14\n\x0charness_name\x18\x03 \x01(\t\x12\x0f\n\x07task_id\x18\x04 \x01(\t\"p\n\x05\x43rash\x12#\n\x06target\x18\x01 \x01(\x0b\x32\x13.msgpb2.BuildOutput\x12\x14\n\x0charness_name\x18\x02 \x01(\t\x12\x18\n\x10\x63rash_input_path\x18\x03 \x01(\t\x12\x12\n\nstacktrace\x18\x04 \x01(\t\"G\n\x16\x43onfirmedVulnerability\x12\x1c\n\x05\x63rash\x18\x01 \x01(\x0b\x32\r.msgpb2.Crash\x12\x0f\n\x07vuln_id\x18\x02 \x01(\t\"A\n\x05Patch\x12\x0f\n\x07task_id\x18\x01 \x01(\t\x12\x18\n\x10vulnerability_id\x18\x02 \x01(\t\x12\r\n\x05patch\x18\x03 \x01(\t\"\x82\x01\n\x0cIndexRequest\x12\x14\n\x0cpackage_name\x18\x01 \x01(\t\x12\x11\n\tsanitizer\x18\x02 \x01(\t\x12\x0f\n\x07ossfuzz\x18\x03 \x01(\t\x12\x13\n\x0bsource_path\x18\x04 \x01(\t\x12\x0f\n\x07task_id\x18\x05 \x01(\t\x12\x12\n\nbuild_type\x18\x06 \x01(\t\"\x81\x01\n\x0bIndexOutput\x12\x14\n\x0cpackage_name\x18\x01 \x01(\t\x12\x11\n\tsanitizer\x18\x02 \x01(\t\x12\x0f\n\x07ossfuzz\x18\x03 \x01(\t\x12\x13\n\x0bsource_path\x18\x04 \x01(\t\x12\x0f\n\x07task_id\x18\x05 \x01(\t\x12\x12\n\nbuild_type\x18\x06 \x01(\tb\x06proto3')
_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, globals())
_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'msg_pb2', globals())
@@ -46,4 +46,8 @@ if _descriptor._USE_C_DESCRIPTORS == False:
_CONFIRMEDVULNERABILITY._serialized_end=1163
_PATCH._serialized_start=1165
_PATCH._serialized_end=1230
_INDEXREQUEST._serialized_start=1233
_INDEXREQUEST._serialized_end=1363
_INDEXOUTPUT._serialized_start=1366
_INDEXOUTPUT._serialized_end=1495
# @@protoc_insertion_point(module_scope)
common/src/buttercup/common/datastructures/msg_pb2.pyi
+32
View File
@@ -62,6 +62,38 @@ class Crash(_message.Message):
target: BuildOutput
def __init__(self, target: _Optional[_Union[BuildOutput, _Mapping]] = ..., harness_name: _Optional[str] = ..., crash_input_path: _Optional[str] = ..., stacktrace: _Optional[str] = ...) -> None: ...
class IndexOutput(_message.Message):
__slots__ = ["build_type", "ossfuzz", "package_name", "sanitizer", "source_path", "task_id"]
BUILD_TYPE_FIELD_NUMBER: _ClassVar[int]
OSSFUZZ_FIELD_NUMBER: _ClassVar[int]
PACKAGE_NAME_FIELD_NUMBER: _ClassVar[int]
SANITIZER_FIELD_NUMBER: _ClassVar[int]
SOURCE_PATH_FIELD_NUMBER: _ClassVar[int]
TASK_ID_FIELD_NUMBER: _ClassVar[int]
build_type: str
ossfuzz: str
package_name: str
sanitizer: str
source_path: str
task_id: str
def __init__(self, package_name: _Optional[str] = ..., sanitizer: _Optional[str] = ..., ossfuzz: _Optional[str] = ..., source_path: _Optional[str] = ..., task_id: _Optional[str] = ..., build_type: _Optional[str] = ...) -> None: ...
class IndexRequest(_message.Message):
__slots__ = ["build_type", "ossfuzz", "package_name", "sanitizer", "source_path", "task_id"]
BUILD_TYPE_FIELD_NUMBER: _ClassVar[int]
OSSFUZZ_FIELD_NUMBER: _ClassVar[int]
PACKAGE_NAME_FIELD_NUMBER: _ClassVar[int]
SANITIZER_FIELD_NUMBER: _ClassVar[int]
SOURCE_PATH_FIELD_NUMBER: _ClassVar[int]
TASK_ID_FIELD_NUMBER: _ClassVar[int]
build_type: str
ossfuzz: str
package_name: str
sanitizer: str
source_path: str
task_id: str
def __init__(self, package_name: _Optional[str] = ..., sanitizer: _Optional[str] = ..., ossfuzz: _Optional[str] = ..., source_path: _Optional[str] = ..., task_id: _Optional[str] = ..., build_type: _Optional[str] = ...) -> None: ...
class Patch(_message.Message):
__slots__ = ["patch", "task_id", "vulnerability_id"]
PATCH_FIELD_NUMBER: _ClassVar[int]
common/src/buttercup/common/queues.py
+29 -1
View File
@@ -13,6 +13,8 @@ from buttercup.common.datastructures.msg_pb2 import (
TaskDelete,
Patch,
ConfirmedVulnerability,
IndexRequest,
IndexOutput,
)
import logging
from typing import Type, Generic, TypeVar, Literal, overload
@@ -31,6 +33,8 @@ class QueueNames(str, Enum):
READY_TASKS = "tasks_ready_queue"
DELETE_TASK = "orchestrator_delete_task_queue"
PATCHES = "patches_queue"
INDEX = "index_queue"
INDEX_OUTPUT = "index_output_queue"
class GroupNames(str, Enum):
@@ -42,6 +46,7 @@ class GroupNames(str, Enum):
SCHEDULER_BUILD_OUTPUT = "scheduler_build_output_group"
UNIQUE_VULNERABILITIES = "unique_vulnerabilities_group"
CONFIRMED_VULNERABILITIES = "confirmed_vulnerabilities_group"
INDEX = "index_group"
class HashNames(str, Enum):
@@ -57,7 +62,8 @@ CRASH_TASK_TIMEOUT_MS = 10 * 60 * 1000
PATCH_TASK_TIMEOUT_MS = 10 * 60 * 1000
UNIQUE_VULNERABILITIES_TASK_TIMEOUT_MS = 10 * 60 * 1000
CONFIRMED_VULNERABILITIES_TASK_TIMEOUT_MS = 10 * 60 * 1000
INDEX_TASK_TIMEOUT_MS = 10 * 60 * 1000
INDEX_OUTPUT_TASK_TIMEOUT_MS = 10 * 60 * 1000
logger = logging.getLogger(__name__)
@@ -312,6 +318,18 @@ class QueueFactory:
PATCH_TASK_TIMEOUT_MS,
[GroupNames.ORCHESTRATOR],
),
QueueNames.INDEX: QueueConfig(
QueueNames.INDEX,
IndexRequest,
INDEX_TASK_TIMEOUT_MS,
[GroupNames.INDEX],
),
QueueNames.INDEX_OUTPUT: QueueConfig(
QueueNames.INDEX_OUTPUT,
IndexOutput,
INDEX_OUTPUT_TASK_TIMEOUT_MS,
[GroupNames.INDEX],
),
}
)
@@ -360,6 +378,16 @@ class QueueFactory:
self, queue_name: Literal[QueueNames.PATCHES], group_name: GroupNames, **kwargs: Any
) -> ReliableQueue[Patch]: ...
@overload
def create(
self, queue_name: Literal[QueueNames.INDEX], group_name: GroupNames, **kwargs: Any
) -> ReliableQueue[IndexRequest]: ...
@overload
def create(
self, queue_name: Literal[QueueNames.INDEX_OUTPUT], group_name: GroupNames, **kwargs: Any
) -> ReliableQueue[IndexOutput]: ...
def create(
self, queue_name: QueueNames, group_name: GroupNames | None = None, **kwargs: Any
) -> ReliableQueue[MsgType]:
compose.yaml
+46 -17
View File
@@ -31,12 +31,6 @@ services:
timeout: 3s
retries: 5
graphdb:
image: janusgraph/janusgraph:1.2.0-20250113-082101.7f6ffcf
restart: always
ports:
- '127.0.0.1:8182:8182'
stimulate-fuzzer-test:
build:
context: ./
@@ -73,27 +67,62 @@ services:
redis:
condition: service_healthy
indexer-run:
image: program-model-image
command: -m buttercup.program_model.indexer.oss_fuzz_indexer --scriptdir /app/scripts --url ws://graphdb:8182/gremlin --allow_pull --oss_fuzz_dir /crs_scratch/ossfuzz --base_image_url gcr.io/oss-fuzz --package_name libpng --wdir /crs_scratch/ --kythe_dir /app/kythe-v0.0.67 --python /app/.venv/bin/python
build:
context: ./
dockerfile: ./program-model/upload_worker.Dockerfile
profiles:
- indexer-run
graphdb:
image: janusgraph/janusgraph:1.2.0-20250126-150810.46a971a
environment:
- janusgraph.storage.backend=cql
- janusgraph.storage.hostname=jg-cassandra
- janusgraph.storage.cql.keyspace=janusgraph
- _JAVA_OPTIONS=-Xms4g -Xmx6g
healthcheck:
test: ["CMD-SHELL", "bin/gremlin.sh", "-e", "scripts/remote-connect.groovy"]
interval: 10s
timeout: 60s
retries: 4
volumes:
- ./crs_scratch:/crs_scratch
- ./program-model/conf/janusgraph-server.yaml:/opt/janusgraph/conf/janusgraph-server.yaml
- ./program-model/load_data.groovy:/opt/janusgraph/load_data.groovy
depends_on:
graphdb-storage:
condition: service_healthy
graphdb-storage:
image: cassandra:4.0.6
container_name: jg-cassandra
environment:
# These values will be modified automatically at competition time
- CASSANDRA_START_RPC=true
healthcheck:
test: ["CMD-SHELL", "nodetool", "status"]
interval: 10s
timeout: 60s
retries: 4
program-model:
build:
context: ./
dockerfile: ./program-model/Dockerfile
command: [
"buttercup-program-model",
"--log_level", "debug",
"serve",
"--wdir", "/crs_scratch",
"--sleep_time", "5",
"--redis_url", "redis://redis:6379",
]
environment:
- BUTTERCUP_PROGRAM_MODEL_MOCK_MODE=true
- DOCKER_HOST=tcp://dind:2375
restart: no
volumes:
- ./crs_scratch:/crs_scratch
- ./tasks_storage:/tasks_storage
depends_on:
redis:
condition: service_healthy
dind:
condition: service_started
graphdb:
condition: service_started
condition: service_healthy
coverage-bot:
command: ["buttercup-coverage-bot", "--wdir", "/crs_scratch", "--redis_url", "redis://redis:6379"]
justfile
+2 -1
View File
@@ -12,7 +12,8 @@ download-kythe:
set -euxo pipefail
if [[ ! -f "program-model/scripts/gzs/kythe-v0.0.67.tar.gz" ]]
then
curl -o program-model/scripts/gzs/kythe-v0.0.67.tar.gz https://github.com/trailofbits/aixcc-kythe/releases/download/v0.0.2/kythe-v0.0.67.tar.gz
mkdir -p program-model/scripts/gzs/
gh release download v0.0.2 -R github.com/trailofbits/aixcc-kythe -D program-model/scripts/gzs/
fi
build-indexer-image:
program-model/.gitignore
+1 -3
View File
@@ -1,5 +1,3 @@
scripts/gzs
*.pyc
program_model/data/kythe/proto
program_model/services/kythe/proto
kythe_builder/build_out
kythe_builder/build_out
program-model/upload_worker.Dockerfile → program-model/Dockerfile
+12 -17
View File
@@ -14,33 +14,28 @@ ENV UV_LINK_MODE=copy \
UV_PYTHON_DOWNLOADS=never \
UV_PYTHON=python3.10
WORKDIR /app
WORKDIR /program-model
RUN --mount=type=cache,target=/root/.cache/uv \
--mount=type=bind,source=common/uv.lock,target=/common/uv.lock \
--mount=type=bind,source=common/pyproject.toml,target=/common/pyproject.toml \
--mount=type=bind,source=common/README.md,target=/common/README.md \
--mount=type=bind,source=program-model/uv.lock,target=uv.lock \
--mount=type=bind,source=program-model/pyproject.toml,target=pyproject.toml \
uv sync --frozen --no-install-project --no-editable
--mount=type=bind,source=common/README.md,target=/common/README.md \
--mount=type=bind,source=program-model/uv.lock,target=/program-model/uv.lock \
--mount=type=bind,source=program-model/pyproject.toml,target=/program-model/pyproject.toml \
--mount=type=bind,source=program-model/README.md,target=/program-model/README.md \
cd /program-model && uv sync --frozen --no-install-project --no-editable
ADD ./common /common
COPY program-model/src /app/src
COPY program-model/pyproject.toml /app/pyproject.toml
COPY program-model/uv.lock /app/uv.lock
COPY program-model/scripts /app/scripts
COPY ./common /common
COPY ./program-model /program-model
RUN --mount=type=cache,target=/root/.cache/uv \
uv sync --frozen --no-editable
FROM base AS runtime
WORKDIR /app
WORKDIR /program-model
COPY --from=builder --chown=app:app /app/scripts /app/scripts
RUN tar -xvf /app/scripts/gzs/kythe-v0.0.67.tar.gz
COPY --from=builder --chown=app:app /app/.venv /app/.venv
COPY ./program-model/scripts /program-model/scripts
ENTRYPOINT ["/app/.venv/bin/python"]
COPY --from=builder --chown=app:app /program-model/.venv /program-model/.venv
ENV PATH=/program-model/.venv/bin:$PATH
program-model/README.md
+158 -41
View File
@@ -1,53 +1,170 @@
# Program Model
This is a demo project that shows how to index an oss-fuzz project (reliablyish) with kythe and upload the results in a graph db
Indexes a program into a graph database.
The components are listed and discussed below:
## Requirements
## kythe_builder
`n2-highmem-8` [instance](https://cloud.google.com/compute/docs/general-purpose-machines#n2-high-mem).
This is a docker image responsible for building kythe within the base-clang image. This should be built first and produces a tar.gz of kythe that can be injected in oss-fuzz docker images.
## Setup
## scripts
Set up working directory and tasks directory.
These scripts are used to create a Docker image that is dervived from a project builder e.g. gcr.io/oss-fuzz/libpng.
```shell
cd afc-crs-trail-of-bits/
The Dockerfile is parameterized over the base image used. So one can build an oss-fuzz project with kythe as follows:
* use oss-fuzz helper.py to build the project
* build the scripts/Dockerfile with the project as the base
cp env.template .env
`cd scripts && docker build --build-arg BASE_IMAGE=gcr.io/oss-fuzz/libpng .`
Then you should be able to run compile as normal for an oss-fuzz image and generate kzips (indexed versions of the object files).
sudo mkdir /crs_scratch/ && sudo chown `whoami`:`whoami` /crs_scratch && sudo mount --bind ./crs_scratch /crs_scratch
## kythe
Once you have a kzip for each compilation action then you need to merge these to a global kzip using
`kythe/tools/kzip merge --output $KYTHE_OUTPUT_DIRECTORY/merged.kzip $KYTHE_OUTPUT_DIRECTORY/*.kzip`
Once you have a bunch of kzips you need to index these with the cxx_indexer in kythe.
e.g.
`kythe/indexers/cxx_indexer --ignore_unimplemented <merged.kzip>`
This produces a stream of protobuf on stdout. Specifically each protobuf record is an entry described in this spec [text](https://kythe.io/docs/kythe-storage.html#_entry)
## program_model
This pyhon script specifically program_model/indexer/entries_into_db.py reads the protobuf entries from kythe and uploads them to JanusGraph.
This is invoked (after entering poetry shell) with
`python program_model/indexer/entries_into_db --url <janus_graph_url>`
janus graph can be run via docker with
`docker run janusgraph/janusgraph`
The url can be determined with
```
docker inspect \
-f '{{range.NetworkSettings.Networks}}{{.IPAddress}}{{end}}' <janus container>
sudo mkdir /tasks_storage && sudo chown `whoami`:`whoami` /tasks_storage && sudo mount --bind ./tasks_storage /tasks_storage
```
this will get you the ip, then the url is `ws://<IP address>:8182/gremlin`
Set up GitHub token.
```shell
gh auth login
GitHub.com
SSH
No
Paste an authentication token
echo <token> | docker login ghcr.io -u <username> --password-stdin
```
Download Kythe.
```shell
cd afc-crs-trail-of-bits/
mkdir -p program-model/scripts/gzs/kythe/
gh release download v0.0.2 -R github.com/trailofbits/aixcc-kythe -D program-model/scripts/gzs/
tar -xvzf program-model/scripts/gzs/kythe-v0.0.67.tar.gz -C program-model/scripts/gzs/kythe/ --strip-components=1
```
## Usage
Prepare example challenge project.
```shell
cd afc-crs-trail-of-bits/
mkdir -p tasks_storage/example-libpng/src/ && mkdir -p tasks_storage/example-libpng/fuzz-tooling/
git clone --branch aixcc-exemplar-challenge-01 git@github.com:aixcc-finals/example-libpng.git tasks_storage/example-libpng/src/example-libpng
git clone --branch master git@github.com:aixcc-finals/oss-fuzz-aixcc.git tasks_storage/example-libpng/fuzz-tooling/fuzz-tooling
```
Start up CRS.
```shell
cd afc-crs-trail-of-bits/
docker compose up -d --build
# cd common/
# uv run src/buttercup/common/msg_publisher.py send tasks_ready_queue ../program-model/mock_data/ready_msg.json
```
Send task to Program Model via Redis queue.
```shell
docker compose logs -f program-model
```
```shell
cd afc-crs-trail-of-bits/program-model/
uv run mock/trigger_pm.py \
--package_name libpng \
--sanitizer AddressSanitizer \
--source_path ../tasks_storage/example-libpng/src/example-libpng/ \
--ossfuzz ../tasks_storage/example-libpng/fuzz-tooling/fuzz-tooling/ \
--task_id example-libpng \
--build_type full
```
Send task to Program Model via cli.
```shell
cd afc-crs-trail-of-bits/program-model/
uv run buttercup-program-model \
--log_level debug \
process \
--wdir ../crs_scratch/ \
--script_dir scripts/ \
--kythe_dir scripts/gzs/kythe/ \
--package_name libpng \
--sanitizer AddressSanitizer \
--source_path ../tasks_storage/example-libpng/src/example-libpng/ \
--ossfuzz ../tasks_storage/example-libpng/fuzz-tooling/fuzz-tooling/ \
--task_id example-libpng \
--build_type full
```
## Testing on Challenges
See [challenges.md](challenges.md).
## Development
Create a new virtual environment and enter it:
```shell
cd afc-crs-trail-of-bits/program-model/
uv venv
source .venv/bin/activate
uv sync --all-extras
```
Run tests:
```shell
pytest
```
Lock, reformat, and lint before committing changes:
```shell
just lock
just reformat
just lint
```
## FAQs
* Why does this use a ubuntu base image?
* Since Kythe uses OSS Fuzz to build and index the challenge source code, we have to use the same base image as ClusterFuzz.
* Why does this use Python 3.10?
* ClusterFuzz uses Python 3.10.
* How do I build Kythe?
* Follow the instructions in [dev.md](dev.md).
## TODO
* [x] Configure Dockerfiles
* [x] Modify `cxxwrapper.sh` to match `ccwrapper.sh`
* [x] Untruncate graph (see `graph.py`)
* [x] Speed up graph creation (see `program_model.py`)
* [ ] Create getter/setter APIs
* [ ] Create unit and regression tests
* [ ] Create one graph per task id
* [ ] Integrate Program Model with other CRS components
* [ ] Coverage Tracker
* [ ] Seed Generator
* [ ] Patcher
* [ ] Verify graph creation correctness
* [ ] Implement more error handling, especially for subprocess runs
* [ ] Download Kythe `.tar.gz` file automatically from github releases
* [ ] Debug issue in `oss_fuzz_indexer.py:index_target()` -- outputted `kzip` files contain critical errors
* [ ] Ensure challenge project is being built and not the oss-fuzz project
* [ ] Periodically backup the graph database into a graphml file
* [ ] Figure out how to make `ccwrapper.sh` not break LD detection
* [ ] Consider periodically removing indexing images for disk space
* [ ] Add support for `java`
* [ ] Consider corner cases
program-model/challenges.md
+28
View File
@@ -0,0 +1,28 @@
# Challenges
Details of challenges tested.
## [libpng](https://github.com/aixcc-finals/example-libpng)
* 2 minutes to write graphml file
* Takes about 8 GB of memory
* Outputted:
* Nodes: 424,350
* Edges: 3,706,932
* 14 minutes to load into JanusGraph
* Takes about 10 GB of memory
* Outputted:
* Nodes: 217,060
* Edges: 1,853,466
**TODO:** Reconcile the differences in the number of nodes and edges.
### Full-scan
* **base:** `2c894c66108f0724331a9e5b4826e351bf2d094b`
### Delta-scan
* **base:** `0cc367aaeaac3f888f255cee5d394968996f736e`
* **delta:** `2c894c66108f0724331a9e5b4826e351bf2d094b`
program-model/conf/janusgraph-server.yaml
+44
View File
@@ -0,0 +1,44 @@
host: 0.0.0.0
port: 8182
evaluationTimeout: 30000
channelizer: org.apache.tinkerpop.gremlin.server.channel.WebSocketChannelizer
graphManager: org.janusgraph.graphdb.management.JanusGraphManager
graphs: {
#graph: conf/janusgraph-inmemory.properties
graph: conf/janusgraph-cql.properties
}
scriptEngines: {
gremlin-groovy: {
plugins: { org.janusgraph.graphdb.tinkerpop.plugin.JanusGraphGremlinPlugin: {},
org.apache.tinkerpop.gremlin.server.jsr223.GremlinServerGremlinPlugin: {},
org.apache.tinkerpop.gremlin.tinkergraph.jsr223.TinkerGraphGremlinPlugin: {},
org.apache.tinkerpop.gremlin.jsr223.ImportGremlinPlugin: {classImports: [java.lang.Math], methodImports: [java.lang.Math#*]},
org.apache.tinkerpop.gremlin.jsr223.ScriptFileGremlinPlugin: {files: [scripts/empty-sample.groovy]}
}
}
}
# serializers:
# - { className: org.apache.tinkerpop.gremlin.util.ser.GraphBinaryMessageSerializerV1, config: { ioRegistries: [org.janusgraph.graphdb.tinkerpop.JanusGraphIoRegistry] }}
# - { className: org.apache.tinkerpop.gremlin.util.ser.GraphBinaryMessageSerializerV1, config: { serializeResultToString: true }}
# - { className: org.apache.tinkerpop.gremlin.util.ser.GraphSONMessageSerializerV3, config: { ioRegistries: [org.janusgraph.graphdb.tinkerpop.JanusGraphIoRegistry] }}
processors:
- { className: org.apache.tinkerpop.gremlin.server.op.session.SessionOpProcessor, config: { sessionTimeout: 28800000 }}
- { className: org.apache.tinkerpop.gremlin.server.op.traversal.TraversalOpProcessor, config: { cacheExpirationTime: 600000, cacheMaxSize: 1000 }}
metrics: {
consoleReporter: {enabled: true, interval: 180000},
csvReporter: {enabled: true, interval: 180000, fileName: /tmp/gremlin-server-metrics.csv},
jmxReporter: {enabled: true},
slf4jReporter: {enabled: true, interval: 180000},
graphiteReporter: {enabled: false, interval: 180000}}
maxInitialLineLength: 4096
maxHeaderSize: 8192
maxChunkSize: 8192
# NOTE(Evan): We need to increase the maxContentLength to 2MB to allow for larger Kythe entries.
# 2MB = 2 * 1024 * 1024
maxContentLength: 2097152
maxAccumulationBufferComponents: 1024
resultIterationBatchSize: 64
writeBufferLowWaterMark: 32768
writeBufferHighWaterMark: 65536
program-model/dev.md
+229
View File
@@ -0,0 +1,229 @@
# Dev
Documentation to get started on developing Program Model APIs.
## Requirements
`n2-standard-4` [instance](https://cloud.google.com/compute/docs/general-purpose-machines#n2_machine_types).
## Docker (WIP)
```shell
(type -p wget >/dev/null || (sudo apt update && sudo apt-get install wget -y)) \
&& sudo mkdir -p -m 755 /etc/apt/keyrings \
&& out=$(mktemp) && wget -nv -O$out https://cli.github.com/packages/githubcli-archive-keyring.gpg \
&& cat $out | sudo tee /etc/apt/keyrings/githubcli-archive-keyring.gpg > /dev/null \
&& sudo chmod go+r /etc/apt/keyrings/githubcli-archive-keyring.gpg \
&& echo "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/githubcli-archive-keyring.gpg] https://cli.github.com/packages stable main" | sudo tee /etc/apt/sources.list.d/github-cli.list > /dev/null \
&& sudo apt update \
&& sudo apt install gh -y
sudo apt update
sudo apt install gh
github.com -> settings -> developer settings -> personal access tokens -> tokens (classic) -> generate new token
select "repo" scope
select "read:org" scope
gh auth login
GitHub.com
SSH
No
Paste an authentication token
gh release download v0.0.2 -R github.com/trailofbits/aixcc-kythe
<authenticate URL in browser for the first time>
rm kythe-v0.0.67.tar.gz
sudo apt install just
cd afc-crs-trail-of-bits/
cp env.template .env
mkdir crs_scratch/
git clone --recursive git@github.com:aixcc-finals/example-libpng.git crs_scratch/libpng
just run-indexer
<todo>
docker compose down
```
## Local
### Setup
```shell
cd afc-crs-trail-of-bits/
cd program-model/
mkdir crs_scratch/
```
### Download libpng
```shell
git clone git@github.com:aixcc-finals/example-libpng.git crs_scratch/libpng
```
### Download Kythe
```shell
mkdir -p scripts/gzs/
gh release download v0.0.2 -R github.com/trailofbits/aixcc-kythe -D scripts/gzs/
mkdir crs_scratch/opt/
tar -zxf scripts/gzs/kythe-v0.0.67.tar.gz --directory crs_scratch/opt/
```
### If no download available, build Kythe
Takes about 3 hours to build.
#### Install Bazel
Follow [instructions](https://bazel.build/install/ubuntu#install-on-ubuntu). You will need to install a specific version of Bazel depending on what the `build` output below outputs.
#### Install Dependencies
```shell
sudo apt install flex bison asciidoc graphviz source-highlight clang
```
#### Download and build Kythe
From [documentation](https://kythe.io/getting-started/#build-a-release-of-kythe-using-bazel-and-unpack-it-in-optkythe)
```shell
git clone git@github.com:trailofbits/aixcc-kythe.git crs_scratch/aixcc-kythe
cd crs_scratch/aixcc-kythe
git checkout a301676c20db9849a06878fa9cc017907eb64d72
bazel build //kythe/release
mkdir ../opt/
tar -zxf bazel-bin/kythe/release/kythe-v0.0.67.tar.gz --directory ../opt/
cd ../
```
## Usage
### Run indexer with Kythe
From [documentation](https://kythe.io/examples/#extracting-cmake-based-repositories)
```shell
sudo apt install cmake
cd crs_scratch/
export KYTHE_ROOT_DIRECTORY=`pwd`/libpng
export KYTHE_CORPUS="myrepo"
export KYTHE_OUTPUT_DIRECTORY=`pwd`/kythe_output
mkdir -p "$KYTHE_OUTPUT_DIRECTORY"
export CMAKE_ROOT_DIRECTORY=`pwd`/libpng
cd libpng/
../opt/kythe-v0.0.67/tools/runextractor cmake -extractor=../opt/kythe-v0.0.67/extractors/cxx_extractor -sourcedir=$CMAKE_ROOT_DIRECTORY
```
### Merge kzip files
```shell
cd crs_scratch/
./opt/kythe-v0.0.67/tools/kzip merge --output $KYTHE_OUTPUT_DIRECTORY/merged.kzip $KYTHE_OUTPUT_DIRECTORY/*.kzip
./opt/kythe-v0.0.67/tools/kzip info --input $KYTHE_OUTPUT_DIRECTORY/merged.kzip | jq .
./opt/kythe-v0.0.67/tools/kzip view $KYTHE_OUTPUT_DIRECTORY/merged.kzip | jq .
```
### Upload to JanusGraph
Extract graphml file from kzip contents.
```shell
cd program-model/src/buttercup
mkdir ../../data/
time ../../crs_scratch/opt/kythe-v0.0.67/indexers/cxx_indexer --ignore_unimplemented ../../crs_scratch/kythe_output/merged.kzip | uv run program_model/indexer/entries_into_graphml.py --output ../../data/graph.xml
```
Create network for JanusGraph to talk to Cassandra backend storage.
```shell
docker network create janusgraph-net
```
Start Cassandra [container](https://github.com/JanusGraph/janusgraph/blob/master/docs/storage-backend/cassandra.md#local-container-mode).
```shell
cd program-model/
./start_cql.sh
```
Start JanusGraph container with custom configuration file.
```shell
cd program-model/
./start_jg.sh
```
Run graph creation.
```shell
cd program-model/
docker ps | grep janusgraph
time docker run --rm --link <container_name>:janusgraph --network janusgraph-net -e GREMLIN_REMOTE_HOSTS=janusgraph -v `pwd`/load_data.groovy:/opt/janusgraph/load_data.groovy janusgraph/janusgraph -- ./bin/gremlin.sh -e load_data.groovy
```
Verify graph creation. From [documentation](https://docs.janusgraph.org/getting-started/installation/).
```shell
docker run --rm --link <container_name>:janusgraph --network janusgraph-net -e GREMLIN_REMOTE_HOSTS=janusgraph -it janusgraph/janusgraph ./bin/gremlin.sh
```
```shell
gremlin> :remote connect tinkerpop.server conf/remote.yaml
gremlin> :remote console
gremlin> g.V().count()
gremlin> g.E().count()
```
Clear graph.
```shell
gremlin> g.V().drop().iterate()
gremlin> g.E().drop().iterate()
```
## Contributing
Before committing code, run the following to ensure that the code is formatted correctly.
```shell
cd program-model/
just reformat
just lint
```
## TODO
* [x] Use Indexes: `org.janusgraph.graphdb.transaction.StandardJanusGraphTx$3.execute - Query requires iterating over all vertices [[~label = x]]. For better performance, use indexes`
* [x] Consider [bulk loading](https://docs.janusgraph.org/operations/bulk-loading/) if needed
* Not required at the moment.
* [x] Consider loading data from a [file](https://tinkerpop.apache.org/docs/3.7.3/dev/io/).
* [ ] Create script to verify creation of graph is accurate.
* [ ] Reconcile the differences in the number of nodes and edges for libpng example.
program-model/justfile
+11
View File
@@ -0,0 +1,11 @@
lock:
uv lock
uv sync --all-extras
lint:
uv run ruff format --check
uv run ruff check
reformat:
uv run ruff format
uv run ruff check --fix
program-model/load_data.groovy
+16
View File
@@ -0,0 +1,16 @@
// We need to configure the console to have state:
// https://tinkerpop.apache.org/docs/current/reference/#console-sessions
:remote connect tinkerpop.server conf/remote.yaml session
:remote console
// 1 hour in milliseconds
:remote config timeout 3600000
println("Starting graph creation...")
g.io('/crs_scratch/graph.xml').read().iterate()
println("Graph creation complete.")
:remote console
:exit
program-model/mock/ready_msg.json
+9
View File
@@ -0,0 +1,9 @@
task {
message_id: "my-message-id"
message_time: 123
task_type: TASK_TYPE_FULL
task_id: "example-libpng"
sources {
url: "https://github.com/buttercup-project/buttercup"
}
}
program-model/mock/trigger_pm.py
+32
View File
@@ -0,0 +1,32 @@
import argparse
from redis import Redis
from buttercup.common.queues import QueueFactory, QueueNames
from buttercup.common.datastructures.msg_pb2 import IndexRequest
def main():
prsr = argparse.ArgumentParser("trigger Program Model manually")
prsr.add_argument("--redis_url", default="redis://127.0.0.1:6379")
prsr.add_argument("--package_name", required=True)
prsr.add_argument("--sanitizer", required=True)
prsr.add_argument("--ossfuzz", required=True)
prsr.add_argument("--source_path", required=True)
prsr.add_argument("--task_id", required=True)
prsr.add_argument("--build_type", required=True)
args = prsr.parse_args()
redis = Redis.from_url(args.redis_url)
queue = QueueFactory(redis).create(QueueNames.INDEX)
req = IndexRequest(
package_name=args.package_name,
sanitizer=args.sanitizer,
ossfuzz=args.ossfuzz,
source_path=args.source_path,
task_id=args.task_id,
build_type=args.build_type,
)
queue.push(req)
if __name__ == "__main__":
main()
program-model/pyproject.toml
+5 -1
View File
@@ -1,7 +1,8 @@
[project]
name = "program-model"
version = "0.1.0"
description = ""
description = "Buttercup contextualizer"
readme = "README.md"
authors = [
{name = "Trail of Bits",email = "opensource@trailofbits"}
]
@@ -12,6 +13,9 @@ dependencies = [
"protobuf (>=3.20, <=3.20.3)",
]
[project.scripts]
buttercup-program-model = "buttercup.program_model.__cli__:main"
[tool.uv.sources]
common = { path = "../common", editable = true }
program-model/scripts/Dockerfile
+11 -14
View File
@@ -1,20 +1,17 @@
ARG BASE_IMAGE=myorg/myapp:latest
FROM $BASE_IMAGE
ARG KYTHE_VERSION=v0.0.67
ARG KYTHE_TAR=gzs/kythe-v0.0.67.tar.gz
FROM ${BASE_IMAGE}
COPY ./ccwrapper.sh $SRC/ccwrapper.sh
COPY ./cxxwrapper.sh $SRC/cxxwrapper.sh
RUN chmod +x $SRC/ccwrapper.sh && chmod +x $SRC/cxxwrapper.sh
COPY ./ccwrapper.sh ${SRC}/ccwrapper.sh
COPY ./cxxwrapper.sh ${SRC}/cxxwrapper.sh
RUN chmod +x ${SRC}/ccwrapper.sh && chmod +x ${SRC}/cxxwrapper.sh
ENV FUZZING_LANGUAGE=c++
ENV ORIG_CXX=$CXX
ENV ORIG_CC=$CC
ENV CXX=$SRC/cxxwrapper.sh
ENV CCC=$SRC/cxxwrapper.sh
ENV CC=$SRC/ccwrapper.sh
ENV ORIG_CXX=${CXX}
ENV ORIG_CC=${CC}
ENV CXX=${SRC}/cxxwrapper.sh
ENV CCC=${SRC}/cxxwrapper.sh
ENV CC=${SRC}/ccwrapper.sh
ADD $KYTHE_TAR $SRC/kythe
#RUN tar -xvf $SRC/kythe-release.tar.gz -C $SRC
ENV KYTHE_RELEASE_DIR=$SRC/kythe/kythe-$KYTHE_VERSION
COPY gzs/kythe/ ${SRC}/kythe
ENV KYTHE_RELEASE_DIR=${SRC}/kythe/
program-model/scripts/ccwrapper.sh
+1 -1
View File
@@ -1,4 +1,4 @@
# cxx_wrapper.sh
# ccwrapper.sh
#!/bin/bash -e
set -euxo pipefail
program-model/scripts/cxxwrapper.sh
+13 -5
View File
@@ -1,9 +1,17 @@
# cxx_wrapper.sh
# cxxwrapper.sh
#!/bin/bash -e
# TODO(Ian): we need to share the code from ccwrapper.sh that moves kzip files between temp and original directories
set -euxo pipefail
echo "Called with $@"
echo "done"
$KYTHE_RELEASE_DIR/extractors/cxx_extractor "$@" & pid=$!
$ORIG_CXX "$@"
wait "$pid"
KYTHE_OUTPUT_DIRECTORY_ORIG=$KYTHE_OUTPUT_DIRECTORY
KYTHE_OUTPUT_DIRECTORY=$(mktemp -d -t kythe_output_XXXXXXXXXXXXXXX)
mkdir -p $KYTHE_OUTPUT_DIRECTORY
KYTHE_OUTPUT_DIRECTORY=$KYTHE_OUTPUT_DIRECTORY $KYTHE_RELEASE_DIR/extractors/cxx_extractor "$@" & pid=$!
$ORIG_CXX "$@"
wait "$pid"
for X in $KYTHE_OUTPUT_DIRECTORY/*.kzip; do [[ -e $X ]] && mv "$X" $KYTHE_OUTPUT_DIRECTORY_ORIG/; done
rm -r $KYTHE_OUTPUT_DIRECTORY
program-model/src/buttercup/program_model/__cli__.py
+60
View File
@@ -0,0 +1,60 @@
import logging
from buttercup.program_model.program_model import ProgramModel
from buttercup.program_model.settings import (
ProgramModelSettings,
ProgramModelServeCommand,
ProgramModelProcessCommand,
)
from buttercup.common.logger import setup_package_logger
from pydantic_settings import get_subcommand
from buttercup.common.datastructures.msg_pb2 import IndexRequest
from redis import Redis
logger = logging.getLogger(__name__)
def prepare_task(command: ProgramModelProcessCommand) -> IndexRequest:
"""Prepares task for indexing."""
return IndexRequest(
package_name=command.package_name,
sanitizer=command.sanitizer,
ossfuzz=command.ossfuzz,
source_path=command.source_path,
task_id=command.task_id,
build_type=command.build_type,
)
def main():
settings = ProgramModelSettings()
setup_package_logger(__name__, settings.log_level)
command = get_subcommand(settings)
if isinstance(command, ProgramModelServeCommand):
redis = Redis.from_url(command.redis_url, decode_responses=False)
with ProgramModel(
sleep_time=command.sleep_time,
redis=redis,
wdir=command.wdir,
script_dir=command.script_dir,
kythe_dir=command.kythe_dir,
python=command.python,
allow_pull=command.allow_pull,
base_image_url=command.base_image_url,
) as program_model:
program_model.serve()
elif isinstance(command, ProgramModelProcessCommand):
task = prepare_task(command)
with ProgramModel(
wdir=command.wdir,
script_dir=command.script_dir,
kythe_dir=command.kythe_dir,
python=command.python,
allow_pull=command.allow_pull,
base_image_url=command.base_image_url,
) as program_model:
program_model.process_task(task)
if __name__ == "__main__":
main()
program-model/src/buttercup/program_model/__init__.py
+4
View File
@@ -0,0 +1,4 @@
"""Program Model."""
__version__ = "0.0.1"
__module_name__ = "program_model"
program-model/src/buttercup/program_model/data/kythe/proto/storage_pb2.py
+35 -39
View File
@@ -2,7 +2,6 @@
# Generated by the protocol buffer compiler. DO NOT EDIT!
# source: kythe/proto/storage.proto
"""Generated protocol buffer code."""
from google.protobuf.internal import builder as _builder
from google.protobuf import descriptor as _descriptor
from google.protobuf import descriptor_pool as _descriptor_pool
@@ -12,45 +11,42 @@ from google.protobuf import symbol_database as _symbol_database
_sym_db = _symbol_database.Default()
DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(
b'\n\x19kythe/proto/storage.proto\x12\x0bkythe.proto"X\n\x05VName\x12\x11\n\tsignature\x18\x01 \x01(\t\x12\x0e\n\x06\x63orpus\x18\x02 \x01(\t\x12\x0c\n\x04root\x18\x03 \x01(\t\x12\x0c\n\x04path\x18\x04 \x01(\t\x12\x10\n\x08language\x18\x05 \x01(\t"\\\n\tVNameMask\x12\x11\n\tsignature\x18\x01 \x01(\x08\x12\x0e\n\x06\x63orpus\x18\x02 \x01(\x08\x12\x0c\n\x04root\x18\x03 \x01(\x08\x12\x0c\n\x04path\x18\x04 \x01(\x08\x12\x10\n\x08language\x18\x05 \x01(\x08"\x89\x01\n\x05\x45ntry\x12"\n\x06source\x18\x01 \x01(\x0b\x32\x12.kythe.proto.VName\x12\x11\n\tedge_kind\x18\x02 \x01(\t\x12"\n\x06target\x18\x03 \x01(\x0b\x32\x12.kythe.proto.VName\x12\x11\n\tfact_name\x18\x04 \x01(\t\x12\x12\n\nfact_value\x18\x05 \x01(\x0c".\n\x07\x45ntries\x12#\n\x07\x65ntries\x18\x01 \x03(\x0b\x32\x12.kythe.proto.Entry"D\n\x0bReadRequest\x12"\n\x06source\x18\x01 \x01(\x0b\x32\x12.kythe.proto.VName\x12\x11\n\tedge_kind\x18\x02 \x01(\t"\xcc\x01\n\x0cWriteRequest\x12"\n\x06source\x18\x01 \x01(\x0b\x32\x12.kythe.proto.VName\x12\x30\n\x06update\x18\x02 \x03(\x0b\x32 .kythe.proto.WriteRequest.Update\x1a\x66\n\x06Update\x12\x11\n\tedge_kind\x18\x01 \x01(\t\x12"\n\x06target\x18\x02 \x01(\x0b\x32\x12.kythe.proto.VName\x12\x11\n\tfact_name\x18\x03 \x01(\t\x12\x12\n\nfact_value\x18\x04 \x01(\x0c"\x0c\n\nWriteReply"Y\n\x0bScanRequest\x12"\n\x06target\x18\x01 \x01(\x0b\x32\x12.kythe.proto.VName\x12\x11\n\tedge_kind\x18\x02 \x01(\t\x12\x13\n\x0b\x66\x61\x63t_prefix\x18\x03 \x01(\t"-\n\x0c\x43ountRequest\x12\r\n\x05index\x18\x01 \x01(\x03\x12\x0e\n\x06shards\x18\x02 \x01(\x03"\x1d\n\nCountReply\x12\x0f\n\x07\x65ntries\x18\x01 \x01(\x03"-\n\x0cShardRequest\x12\r\n\x05index\x18\x01 \x01(\x03\x12\x0e\n\x06shards\x18\x02 \x01(\x03"N\n\x10VNameRewriteRule\x12\x0f\n\x07pattern\x18\x01 \x01(\t\x12)\n\x06v_name\x18\x02 \x01(\x0b\x32\x12.kythe.proto.VNameR\x05vname"@\n\x11VNameRewriteRules\x12+\n\x04rule\x18\x01 \x03(\x0b\x32\x1d.kythe.proto.VNameRewriteRuleBH\n\x1f\x63om.google.devtools.kythe.protoZ%kythe.io/kythe/proto/storage_go_protob\x06proto3'
)
DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x19kythe/proto/storage.proto\x12\x0bkythe.proto\"X\n\x05VName\x12\x11\n\tsignature\x18\x01 \x01(\t\x12\x0e\n\x06\x63orpus\x18\x02 \x01(\t\x12\x0c\n\x04root\x18\x03 \x01(\t\x12\x0c\n\x04path\x18\x04 \x01(\t\x12\x10\n\x08language\x18\x05 \x01(\t\"\\\n\tVNameMask\x12\x11\n\tsignature\x18\x01 \x01(\x08\x12\x0e\n\x06\x63orpus\x18\x02 \x01(\x08\x12\x0c\n\x04root\x18\x03 \x01(\x08\x12\x0c\n\x04path\x18\x04 \x01(\x08\x12\x10\n\x08language\x18\x05 \x01(\x08\"\x89\x01\n\x05\x45ntry\x12\"\n\x06source\x18\x01 \x01(\x0b\x32\x12.kythe.proto.VName\x12\x11\n\tedge_kind\x18\x02 \x01(\t\x12\"\n\x06target\x18\x03 \x01(\x0b\x32\x12.kythe.proto.VName\x12\x11\n\tfact_name\x18\x04 \x01(\t\x12\x12\n\nfact_value\x18\x05 \x01(\x0c\".\n\x07\x45ntries\x12#\n\x07\x65ntries\x18\x01 \x03(\x0b\x32\x12.kythe.proto.Entry\"D\n\x0bReadRequest\x12\"\n\x06source\x18\x01 \x01(\x0b\x32\x12.kythe.proto.VName\x12\x11\n\tedge_kind\x18\x02 \x01(\t\"\xcc\x01\n\x0cWriteRequest\x12\"\n\x06source\x18\x01 \x01(\x0b\x32\x12.kythe.proto.VName\x12\x30\n\x06update\x18\x02 \x03(\x0b\x32 .kythe.proto.WriteRequest.Update\x1a\x66\n\x06Update\x12\x11\n\tedge_kind\x18\x01 \x01(\t\x12\"\n\x06target\x18\x02 \x01(\x0b\x32\x12.kythe.proto.VName\x12\x11\n\tfact_name\x18\x03 \x01(\t\x12\x12\n\nfact_value\x18\x04 \x01(\x0c\"\x0c\n\nWriteReply\"Y\n\x0bScanRequest\x12\"\n\x06target\x18\x01 \x01(\x0b\x32\x12.kythe.proto.VName\x12\x11\n\tedge_kind\x18\x02 \x01(\t\x12\x13\n\x0b\x66\x61\x63t_prefix\x18\x03 \x01(\t\"-\n\x0c\x43ountRequest\x12\r\n\x05index\x18\x01 \x01(\x03\x12\x0e\n\x06shards\x18\x02 \x01(\x03\"\x1d\n\nCountReply\x12\x0f\n\x07\x65ntries\x18\x01 \x01(\x03\"-\n\x0cShardRequest\x12\r\n\x05index\x18\x01 \x01(\x03\x12\x0e\n\x06shards\x18\x02 \x01(\x03\"N\n\x10VNameRewriteRule\x12\x0f\n\x07pattern\x18\x01 \x01(\t\x12)\n\x06v_name\x18\x02 \x01(\x0b\x32\x12.kythe.proto.VNameR\x05vname\"@\n\x11VNameRewriteRules\x12+\n\x04rule\x18\x01 \x03(\x0b\x32\x1d.kythe.proto.VNameRewriteRuleBH\n\x1f\x63om.google.devtools.kythe.protoZ%kythe.io/kythe/proto/storage_go_protob\x06proto3')
_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, globals())
_builder.BuildTopDescriptorsAndMessages(
DESCRIPTOR, "kythe.proto.storage_pb2", globals()
)
_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'kythe.proto.storage_pb2', globals())
if _descriptor._USE_C_DESCRIPTORS == False:
DESCRIPTOR._options = None
DESCRIPTOR._serialized_options = (
b"\n\037com.google.devtools.kythe.protoZ%kythe.io/kythe/proto/storage_go_proto"
)
_VNAME._serialized_start = 42
_VNAME._serialized_end = 130
_VNAMEMASK._serialized_start = 132
_VNAMEMASK._serialized_end = 224
_ENTRY._serialized_start = 227
_ENTRY._serialized_end = 364
_ENTRIES._serialized_start = 366
_ENTRIES._serialized_end = 412
_READREQUEST._serialized_start = 414
_READREQUEST._serialized_end = 482
_WRITEREQUEST._serialized_start = 485
_WRITEREQUEST._serialized_end = 689
_WRITEREQUEST_UPDATE._serialized_start = 587
_WRITEREQUEST_UPDATE._serialized_end = 689
_WRITEREPLY._serialized_start = 691
_WRITEREPLY._serialized_end = 703
_SCANREQUEST._serialized_start = 705
_SCANREQUEST._serialized_end = 794
_COUNTREQUEST._serialized_start = 796
_COUNTREQUEST._serialized_end = 841
_COUNTREPLY._serialized_start = 843
_COUNTREPLY._serialized_end = 872
_SHARDREQUEST._serialized_start = 874
_SHARDREQUEST._serialized_end = 919
_VNAMEREWRITERULE._serialized_start = 921
_VNAMEREWRITERULE._serialized_end = 999
_VNAMEREWRITERULES._serialized_start = 1001
_VNAMEREWRITERULES._serialized_end = 1065
DESCRIPTOR._options = None
DESCRIPTOR._serialized_options = b'\n\037com.google.devtools.kythe.protoZ%kythe.io/kythe/proto/storage_go_proto'
_VNAME._serialized_start=42
_VNAME._serialized_end=130
_VNAMEMASK._serialized_start=132
_VNAMEMASK._serialized_end=224
_ENTRY._serialized_start=227
_ENTRY._serialized_end=364
_ENTRIES._serialized_start=366
_ENTRIES._serialized_end=412
_READREQUEST._serialized_start=414
_READREQUEST._serialized_end=482
_WRITEREQUEST._serialized_start=485
_WRITEREQUEST._serialized_end=689
_WRITEREQUEST_UPDATE._serialized_start=587
_WRITEREQUEST_UPDATE._serialized_end=689
_WRITEREPLY._serialized_start=691
_WRITEREPLY._serialized_end=703
_SCANREQUEST._serialized_start=705
_SCANREQUEST._serialized_end=794
_COUNTREQUEST._serialized_start=796
_COUNTREQUEST._serialized_end=841
_COUNTREPLY._serialized_start=843
_COUNTREPLY._serialized_end=872
_SHARDREQUEST._serialized_start=874
_SHARDREQUEST._serialized_end=919
_VNAMEREWRITERULE._serialized_start=921
_VNAMEREWRITERULE._serialized_end=999
_VNAMEREWRITERULES._serialized_start=1001
_VNAMEREWRITERULES._serialized_end=1065
# @@protoc_insertion_point(module_scope)
program-model/src/buttercup/program_model/data/kythe/proto/storage_pb2.pyi
+13 -72
View File
@@ -1,13 +1,7 @@
from google.protobuf.internal import containers as _containers
from google.protobuf import descriptor as _descriptor
from google.protobuf import message as _message
from typing import (
ClassVar as _ClassVar,
Iterable as _Iterable,
Mapping as _Mapping,
Optional as _Optional,
Union as _Union,
)
from typing import ClassVar as _ClassVar, Iterable as _Iterable, Mapping as _Mapping, Optional as _Optional, Union as _Union
DESCRIPTOR: _descriptor.FileDescriptor
@@ -23,17 +17,13 @@ class CountRequest(_message.Message):
SHARDS_FIELD_NUMBER: _ClassVar[int]
index: int
shards: int
def __init__(
self, index: _Optional[int] = ..., shards: _Optional[int] = ...
) -> None: ...
def __init__(self, index: _Optional[int] = ..., shards: _Optional[int] = ...) -> None: ...
class Entries(_message.Message):
__slots__ = ["entries"]
ENTRIES_FIELD_NUMBER: _ClassVar[int]
entries: _containers.RepeatedCompositeFieldContainer[Entry]
def __init__(
self, entries: _Optional[_Iterable[_Union[Entry, _Mapping]]] = ...
) -> None: ...
def __init__(self, entries: _Optional[_Iterable[_Union[Entry, _Mapping]]] = ...) -> None: ...
class Entry(_message.Message):
__slots__ = ["edge_kind", "fact_name", "fact_value", "source", "target"]
@@ -47,14 +37,7 @@ class Entry(_message.Message):
fact_value: bytes
source: VName
target: VName
def __init__(
self,
source: _Optional[_Union[VName, _Mapping]] = ...,
edge_kind: _Optional[str] = ...,
target: _Optional[_Union[VName, _Mapping]] = ...,
fact_name: _Optional[str] = ...,
fact_value: _Optional[bytes] = ...,
) -> None: ...
def __init__(self, source: _Optional[_Union[VName, _Mapping]] = ..., edge_kind: _Optional[str] = ..., target: _Optional[_Union[VName, _Mapping]] = ..., fact_name: _Optional[str] = ..., fact_value: _Optional[bytes] = ...) -> None: ...
class ReadRequest(_message.Message):
__slots__ = ["edge_kind", "source"]
@@ -62,11 +45,7 @@ class ReadRequest(_message.Message):
SOURCE_FIELD_NUMBER: _ClassVar[int]
edge_kind: str
source: VName
def __init__(
self,
source: _Optional[_Union[VName, _Mapping]] = ...,
edge_kind: _Optional[str] = ...,
) -> None: ...
def __init__(self, source: _Optional[_Union[VName, _Mapping]] = ..., edge_kind: _Optional[str] = ...) -> None: ...
class ScanRequest(_message.Message):
__slots__ = ["edge_kind", "fact_prefix", "target"]
@@ -76,12 +55,7 @@ class ScanRequest(_message.Message):
edge_kind: str
fact_prefix: str
target: VName
def __init__(
self,
target: _Optional[_Union[VName, _Mapping]] = ...,
edge_kind: _Optional[str] = ...,
fact_prefix: _Optional[str] = ...,
) -> None: ...
def __init__(self, target: _Optional[_Union[VName, _Mapping]] = ..., edge_kind: _Optional[str] = ..., fact_prefix: _Optional[str] = ...) -> None: ...
class ShardRequest(_message.Message):
__slots__ = ["index", "shards"]
@@ -89,9 +63,7 @@ class ShardRequest(_message.Message):
SHARDS_FIELD_NUMBER: _ClassVar[int]
index: int
shards: int
def __init__(
self, index: _Optional[int] = ..., shards: _Optional[int] = ...
) -> None: ...
def __init__(self, index: _Optional[int] = ..., shards: _Optional[int] = ...) -> None: ...
class VName(_message.Message):
__slots__ = ["corpus", "language", "path", "root", "signature"]
@@ -105,14 +77,7 @@ class VName(_message.Message):
path: str
root: str
signature: str
def __init__(
self,
signature: _Optional[str] = ...,
corpus: _Optional[str] = ...,
root: _Optional[str] = ...,
path: _Optional[str] = ...,
language: _Optional[str] = ...,
) -> None: ...
def __init__(self, signature: _Optional[str] = ..., corpus: _Optional[str] = ..., root: _Optional[str] = ..., path: _Optional[str] = ..., language: _Optional[str] = ...) -> None: ...
class VNameMask(_message.Message):
__slots__ = ["corpus", "language", "path", "root", "signature"]
@@ -126,14 +91,7 @@ class VNameMask(_message.Message):
path: bool
root: bool
signature: bool
def __init__(
self,
signature: bool = ...,
corpus: bool = ...,
root: bool = ...,
path: bool = ...,
language: bool = ...,
) -> None: ...
def __init__(self, signature: bool = ..., corpus: bool = ..., root: bool = ..., path: bool = ..., language: bool = ...) -> None: ...
class VNameRewriteRule(_message.Message):
__slots__ = ["pattern", "v_name"]
@@ -141,19 +99,13 @@ class VNameRewriteRule(_message.Message):
V_NAME_FIELD_NUMBER: _ClassVar[int]
pattern: str
v_name: VName
def __init__(
self,
pattern: _Optional[str] = ...,
v_name: _Optional[_Union[VName, _Mapping]] = ...,
) -> None: ...
def __init__(self, pattern: _Optional[str] = ..., v_name: _Optional[_Union[VName, _Mapping]] = ...) -> None: ...
class VNameRewriteRules(_message.Message):
__slots__ = ["rule"]
RULE_FIELD_NUMBER: _ClassVar[int]
rule: _containers.RepeatedCompositeFieldContainer[VNameRewriteRule]
def __init__(
self, rule: _Optional[_Iterable[_Union[VNameRewriteRule, _Mapping]]] = ...
) -> None: ...
def __init__(self, rule: _Optional[_Iterable[_Union[VNameRewriteRule, _Mapping]]] = ...) -> None: ...
class WriteReply(_message.Message):
__slots__ = []
@@ -171,20 +123,9 @@ class WriteRequest(_message.Message):
fact_name: str
fact_value: bytes
target: VName
def __init__(
self,
edge_kind: _Optional[str] = ...,
target: _Optional[_Union[VName, _Mapping]] = ...,
fact_name: _Optional[str] = ...,
fact_value: _Optional[bytes] = ...,
) -> None: ...
def __init__(self, edge_kind: _Optional[str] = ..., target: _Optional[_Union[VName, _Mapping]] = ..., fact_name: _Optional[str] = ..., fact_value: _Optional[bytes] = ...) -> None: ...
SOURCE_FIELD_NUMBER: _ClassVar[int]
UPDATE_FIELD_NUMBER: _ClassVar[int]
source: VName
update: _containers.RepeatedCompositeFieldContainer[WriteRequest.Update]
def __init__(
self,
source: _Optional[_Union[VName, _Mapping]] = ...,
update: _Optional[_Iterable[_Union[WriteRequest.Update, _Mapping]]] = ...,
) -> None: ...
def __init__(self, source: _Optional[_Union[VName, _Mapping]] = ..., update: _Optional[_Iterable[_Union[WriteRequest.Update, _Mapping]]] = ...) -> None: ...
program-model/src/buttercup/program_model/graph.py
+209
View File
@@ -0,0 +1,209 @@
import urllib.parse
from xml.dom import minidom
import sys
import argparse
import logging
from dataclasses import dataclass, field
from io import BytesIO
from typing import Generator, Dict, Set
from buttercup.program_model.data.kythe.proto.storage_pb2 import Entry, VName
from buttercup.program_model.utils.varint import decode_stream
import base64
logger = logging.getLogger(__name__)
@dataclass(frozen=True, repr=False)
class KytheURI:
"""Kythe entries described here: https://kythe.io/docs/kythe-storage.html#_entry"""
signature: str
corpus: str
root: str
path: str
language: str
def __str__(self):
lst = [self.signature, self.corpus, self.root, self.path, self.language]
uristr = "/".join([urllib.parse.quote(x, safe="") for x in lst])
return uristr
@staticmethod
def from_vname(vn: VName):
return KytheURI(
signature=vn.signature,
corpus=vn.corpus,
root=vn.root,
path=vn.path,
language=vn.language,
)
@dataclass
class Node:
"""Node in the graph."""
id: str
property: Dict[str, str] = field(default_factory=dict)
def to_graphml(self) -> str:
"""Convert node to a GraphML string."""
content = []
content.append(f'<node id="{self.id}">')
for key, value in self.property.items():
content.append(f'<data key="{key}">{value}</data>')
content.append("</node>")
return "".join(content)
@dataclass
class Edge:
"""Edge in the graph."""
id: str
source_id: str
target_id: str
property: Dict[str, str] = field(default_factory=dict)
def to_graphml(self) -> str:
"""Convert edge to a GraphML string."""
content = []
content.append(
f'<edge id="{self.id}" source="{self.source_id}" target="{self.target_id}">'
)
for key, value in self.property.items():
content.append(f'<data key="{key}">{value}</data>')
content.append("</edge>")
return "".join(content)
@dataclass(repr=False)
class GraphStorage:
"""Class to interact between Kythe and an output file."""
def __init__(self):
self.nodes: Dict[str, Node] = {}
self.node_properties: Set[str] = set()
self.edges: Dict[str, Edge] = {}
self.edge_properties: Set[str] = set()
def is_edge(self, ent: Entry) -> bool:
"""Check if the entry is an edge."""
return ent.edge_kind != ""
def encode_value(self, value: bytes) -> str:
"""Encode a value in base64."""
return base64.b64encode(value).decode()
def convert_node(self, nd: VName) -> Node:
"""Converts a Kythe node to a GraphML node."""
uri = str(KytheURI.from_vname(nd))
if uri in self.nodes.keys():
return self.nodes[uri]
return Node(id=uri)
def process_stream(self, fl: BytesIO):
"""Process a stream of Kythe entries and output them to a GraphML file."""
try:
for count, entry in enumerate(self.iterate_over_entries(fl)):
source_node = self.convert_node(entry.source)
key = entry.fact_name
value = self.encode_value(entry.fact_value)
if self.is_edge(entry):
target_node = self.convert_node(entry.target)
edge = Edge(
id=len(self.edges.keys()),
source_id=source_node.id,
target_id=target_node.id,
)
edge.property["labelE"] = entry.edge_kind
self.edge_properties.add("labelE")
self.edges[edge.id] = edge
else:
source_node.property[key] = value
self.node_properties.add(key)
self.nodes[source_node.id] = source_node
# TODO(Evan): Remove this after diagnosing graph creation error
if count == 1000:
break
except Exception as e:
logger.error("Exception occurred: %s", e)
raise e
def iterate_over_entries(self, fl: BytesIO) -> Generator[Entry, None, None]:
"""Iterate over entries in the stream."""
# Indexers emit a delimited stream of entry protobufs
# From: https://kythe.io/examples/#indexing-compilations
while True:
try:
yield self.parse_entry(fl)
# This is the end of the stream. This error is expected.
except EOFError:
break
except Exception as e:
raise e
def parse_entry(self, fl: BytesIO) -> Entry:
"""Parse a Kythe entry from the stream."""
# Read the size of the entry protobuf
sz = decode_stream(fl)
bts = bytes()
while len(bts) < sz:
rd = fl.read(sz - len(bts))
if len(rd) <= 0:
return None
bts += rd
ent = Entry()
ent.ParseFromString(bts)
return ent
def to_graphml(self) -> str:
"""Convert graph to a GraphML file.
From: https://tinkerpop.apache.org/docs/3.7.3/dev/io/
"""
content = '<?xml version="1.0" encoding="UTF-8"?>'
content += '<graphml xmlns="http://graphml.graphdrawing.org/xmlns" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://graphml.graphdrawing.org/xmlns http://graphml.graphdrawing.org/xmlns/1.1/graphml.xsd">'
content += '<graph id="G" edgedefault="directed">'
# Output node properties
for p in self.node_properties:
content += f'<key id="{p}" for="node" attr.name="{p}" attr.type="string" />'
# Output edge properties
for p in self.edge_properties:
content += f'<key id="{p}" for="edge" attr.name="{p}" attr.type="string" />'
# Output node contents
for node in self.nodes.values():
content += node.to_graphml()
# Output edge contents
for edge in self.edges.values():
content += edge.to_graphml()
content += "</graph>"
content += "</graphml>"
xml_doc = minidom.parseString(content)
pretty_xml = xml_doc.toprettyxml(indent=" ")
return pretty_xml
def main():
prsr = argparse.ArgumentParser("Entry Output to GraphML")
prsr.add_argument("--output", required=True)
args = prsr.parse_args()
with open(args.output, "w") as fw, sys.stdin.buffer as fr:
gs = GraphStorage()
gs.process_stream(fr)
fw.write(gs.to_graphml())
if __name__ == "__main__":
main()
program-model/src/buttercup/program_model/indexer/oss_fuzz_indexer.py → program-model/src/buttercup/program_model/indexer.py
+10 -42
View File
@@ -1,13 +1,15 @@
import argparse
import logging
import uuid
import subprocess
from dataclasses import dataclass
from buttercup.common import oss_fuzz_tool
from buttercup.common.oss_fuzz_tool import OSSFuzzTool
from buttercup.program_model.indexer.entries_into_db import JanusStorage
from buttercup.program_model.kythe import KytheTool, KytheConf
import os
logger = logging.getLogger(__name__)
@dataclass
class IndexTarget:
@@ -16,9 +18,8 @@ class IndexTarget:
@dataclass
class Conf:
class IndexConf:
scriptdir: str
url: str
python: str
allow_pull: bool
base_image_url: str
@@ -26,7 +27,7 @@ class Conf:
class Indexer:
def __init__(self, conf: Conf):
def __init__(self, conf: IndexConf):
self.conf = conf
def build_image(self, idx_target: IndexTarget):
@@ -39,7 +40,7 @@ class Indexer:
)
)
base_image_name = fzz_tool.build_base_image(idx_target.package_name)
print(base_image_name)
logger.debug(f"Base image name: {base_image_name}")
if base_image_name is None:
return None
@@ -85,38 +86,9 @@ class Indexer:
return output_dir
@dataclass
class KytheConf:
kythe_dir: str
class KytheTool:
def __init__(self, conf: KytheConf):
self.conf = conf
def merge_kythe_output(self, input_dir: str, output_kzip: str):
merge_path = os.path.join(self.conf.kythe_dir, "tools/kzip")
total = []
for fl in os.listdir(input_dir):
if fl.endswith(".kzip"):
total.append(os.path.join(input_dir, fl))
command = [merge_path, "merge", "--output", output_kzip] + total
subprocess.run(command, check=True)
return True
def cxx_index(self, input_kzip: str, output_bin: str):
indexer_path = os.path.join(self.conf.kythe_dir, "indexers/cxx_indexer")
command = [indexer_path, "-i", input_kzip, "-o", output_bin]
subprocess.run(command, check=True)
return True
def main():
prsr = argparse.ArgumentParser("oss fuzz builder")
prsr.add_argument("--scriptdir", required=True)
prsr.add_argument("--url", required=True)
prsr.add_argument("--python", default="python")
prsr.add_argument("--allow_pull", action="store_true", default=False)
prsr.add_argument("--base_image_url", required=True)
@@ -126,9 +98,8 @@ def main():
prsr.add_argument("--kythe_dir", required=True)
args = prsr.parse_args()
conf = Conf(
conf = IndexConf(
args.scriptdir,
args.url,
args.python,
args.allow_pull,
args.base_image_url,
@@ -137,9 +108,9 @@ def main():
indexer = Indexer(conf)
output_dir = indexer.index_target(IndexTarget(args.oss_fuzz_dir, args.package_name))
if output_dir is None:
print("Failed to index target")
logger.error("Failed to index target")
return
print(output_dir)
logger.debug(f"Output dir: {output_dir}")
output_id = str(uuid.uuid4())
ktool = KytheTool(KytheConf(args.kythe_dir))
merged_kzip = os.path.join(args.wdir, f"kythe_output_merge_{output_id}.kzip")
@@ -147,9 +118,6 @@ def main():
cxx_bin = os.path.join(args.wdir, f"kythe_output_cxx_{output_id}.bin")
ktool.cxx_index(merged_kzip, cxx_bin)
with open(cxx_bin, "rb") as f:
JanusStorage(args.url).process_stream(args.package_name, f)
if __name__ == "__main__":
main()
program-model/src/buttercup/program_model/indexer/__init__.py
View File
program-model/src/buttercup/program_model/indexer/entries_into_db.py
-110
View File
@@ -1,110 +0,0 @@
from buttercup.program_model.data.kythe.proto.storage_pb2 import Entry, VName
from buttercup.program_model.utils.varint import decode_stream
from typing import Generator
from io import BytesIO
from gremlin_python.structure.graph import Graph, GraphTraversalSource
from gremlin_python.driver.driver_remote_connection import DriverRemoteConnection
from dataclasses import dataclass
import urllib
from gremlin_python.structure.graph import Vertex
import sys
import argparse
@dataclass(frozen=True, repr=False)
class KytheURI:
signature: str
corpus: str
root: str
path: str
language: str
def __str__(self):
lst = [self.signature, self.corpus, self.root, self.path, self.language]
uristr = "/".join([urllib.parse.quote(x, safe="") for x in lst])
return uristr
@staticmethod
def from_vname(vn: VName):
return KytheURI(
signature=vn.signature,
corpus=vn.corpus,
root=vn.root,
path=vn.path,
language=vn.language,
)
def is_edge(ent: Entry) -> bool:
return ent.edge_kind != ""
def convert_node(trv: GraphTraversalSource, nd: VName) -> Vertex:
ur = KytheURI.from_vname(nd)
s = str(ur)
print(s)
maybe_nd = next(trv.V().has_label(s), None)
if maybe_nd is not None:
return maybe_nd
return trv.add_v(str(ur)).next()
class JanusStorage:
def __init__(self, url: str):
self.connection = DriverRemoteConnection(url, traversal_source="g")
self.graph = Graph()
self.g = self.graph.traversal().with_remote(self.connection)
next(self.g.V().has_label("x"), None)
def process_stream(self, _project_name: str, fl: BytesIO):
next(self.g.V().has_label("x"), None)
tx = self.g.tx()
trv = tx.begin()
for ent in iterate_over_entries(fl):
next(trv.V().has_label("x"), None)
if is_edge(ent):
n1 = convert_node(trv, ent.source)
n2 = convert_node(trv, ent.target)
trv.add_e().from_(n1).to(n2).property(
"edge_kind", ent.edge_kind
).iterate()
else:
n1 = convert_node(trv, ent.source)
trv.V().has_id(n1.id).property(
ent.fact_name, ent.fact_value.decode()
).iterate()
tx.commit()
def iterate_over_entries(fl: BytesIO) -> Generator[Entry, None, None]:
try:
sz = decode_stream(fl)
bts = bytes()
while len(bts) < sz:
rd = fl.read(sz - len(bts))
if len(rd) <= 0:
return None
bts += rd
ent = Entry()
ent.ParseFromString(bts)
yield ent
except EOFError:
return None
except Exception:
# TODO(Ian) should catch decode errors
return None
def main():
prsr = argparse.ArgumentParser("Entry Upload to Janus")
prsr.add_argument("--url", required=True)
args = prsr.parse_args()
storage = JanusStorage(args.url)
with sys.stdin.buffer as f:
storage.process_stream("", f)
if __name__ == "__main__":
main()
program-model/src/buttercup/program_model/kythe.py
+34
View File
@@ -0,0 +1,34 @@
import subprocess
import logging
import os
from dataclasses import dataclass
logger = logging.getLogger(__name__)
@dataclass
class KytheConf:
kythe_dir: str
class KytheTool:
def __init__(self, conf: KytheConf):
self.conf = conf
def merge_kythe_output(self, input_dir: str, output_kzip: str):
merge_path = os.path.join(self.conf.kythe_dir, "tools/kzip")
total = []
for fl in os.listdir(input_dir):
if fl.endswith(".kzip"):
total.append(os.path.join(input_dir, fl))
command = [merge_path, "merge", "--output", output_kzip] + total
subprocess.run(command, check=True)
return True
def cxx_index(self, input_kzip: str, output_bin: str):
indexer_path = os.path.join(self.conf.kythe_dir, "indexers/cxx_indexer")
command = [indexer_path, input_kzip, "-o", output_bin]
subprocess.run(command, check=True)
return True
program-model/src/buttercup/program_model/logger.py
+25
View File
@@ -0,0 +1,25 @@
"""Configure loggers."""
from __future__ import annotations
import logging
import buttercup.program_model
MAIN_LOGGER_NAME = buttercup.program_model.__module_name__
def logger_configurer(log_level: str | int) -> None:
"""Configure main logger."""
logger = logging.getLogger(MAIN_LOGGER_NAME)
logger.setLevel(log_level)
handler = logging.StreamHandler()
handler.setLevel(log_level)
formatter = logging.Formatter(
"%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
handler.setFormatter(formatter)
if logger.hasHandlers():
logger.handlers.clear()
logger.addHandler(handler)
logger.propagate = False
program-model/src/buttercup/program_model/program_model.py
+207
View File
@@ -0,0 +1,207 @@
import logging
import os
import stat
import shutil
from dataclasses import dataclass, field
from buttercup.common.queues import (
RQItem,
QueueFactory,
ReliableQueue,
QueueNames,
GroupNames,
)
from buttercup.program_model.indexer import Indexer, IndexConf, IndexTarget
from buttercup.program_model.kythe import KytheTool, KytheConf
from buttercup.program_model.graph import GraphStorage
from buttercup.common.datastructures.msg_pb2 import IndexRequest, IndexOutput
from buttercup.common.challenge_task import ChallengeTask
from pathlib import Path
from redis import Redis
import time
import subprocess
import uuid
import tempfile
from typing import Optional
logger = logging.getLogger(__name__)
@dataclass
class ProgramModel:
sleep_time: float = 1.0
redis: Redis | None = None
task_queue: ReliableQueue | None = field(init=False, default=None)
output_queue: ReliableQueue | None = field(init=False, default=None)
wdir: Path | None = None
script_dir: Path | None = None
kythe_dir: Path | None = None
python: str | None = None
allow_pull: bool = True
base_image_url: str = "gcr.io/oss-fuzz"
def __post_init__(self):
"""Post-initialization setup."""
self.wdir = Path(self.wdir).resolve()
self.script_dir = Path(self.script_dir).resolve()
self.kythe_dir = Path(self.kythe_dir).resolve()
if self.redis is not None:
logger.debug("Using Redis for task queues")
queue_factory = QueueFactory(self.redis)
self.task_queue = queue_factory.create(QueueNames.INDEX, GroupNames.INDEX)
self.output_queue = queue_factory.create(QueueNames.INDEX_OUTPUT)
def __enter__(self):
return self
def __exit__(self, exc_type, exc_val, exc_tb):
self.cleanup()
def cleanup(self):
"""Cleanup resources used by the program model"""
pass
def process_task(self, args: IndexRequest) -> bool:
"""Process a single task for indexing a program"""
# Convert path strings to Path objects
ossfuzz = Path(args.ossfuzz).resolve()
source_path = Path(args.source_path).resolve()
task_dir = source_path.parent.parent.resolve()
logger.debug(f"Kythe dir: {self.kythe_dir}")
logger.debug(f"Script dir: {self.script_dir}")
logger.debug(f"Wdir: {self.wdir}")
logger.debug(f"Task dir: {task_dir}")
logger.debug(f"Source dir: {source_path}")
logger.debug(f"OSSFuzz dir: {ossfuzz}")
logger.debug(f"Python: {self.python}")
logger.debug(f"Allow pull: {self.allow_pull}")
logger.debug(f"Base image URL: {self.base_image_url}")
with tempfile.TemporaryDirectory(dir=self.wdir) as td:
logger.info(
f"Running indexer for {args.package_name} | {source_path} | {args.task_id}"
)
# Change permissions so that JanusGraph can read from the temporary directory
current = os.stat(td).st_mode
janus_user = stat.S_IRGRP | stat.S_IROTH | stat.S_IXGRP | stat.S_IXOTH
os.chmod(td, current | janus_user)
tsk = ChallengeTask(
read_only_task_dir=task_dir,
project_name=args.package_name,
python_path=self.python,
)
with tsk.get_rw_copy(work_dir=td) as local_tsk:
logger.debug(f"Local path: {local_tsk.local_task_dir}")
# Index the program
indexer_conf = IndexConf(
scriptdir=self.script_dir,
python=self.python,
allow_pull=self.allow_pull,
base_image_url=self.base_image_url,
wdir=td,
)
indexer = Indexer(indexer_conf)
output_dir = indexer.index_target(
IndexTarget(
oss_fuzz_dir=ossfuzz,
package_name=args.package_name,
)
)
if output_dir is None:
logger.error(f"Failed to index program {args.package_name}")
return False
else:
logger.info(f"Successfully indexed program {args.package_name}")
# Merge index files
output_id = str(uuid.uuid4())
ktool = KytheTool(KytheConf(self.kythe_dir))
merged_kzip = Path(td) / f"kythe_output_merge_{output_id}.kzip"
ktool.merge_kythe_output(output_dir, merged_kzip)
# Convert the merged kzip file into a binary file
bin_file = Path(td) / f"kythe_output_cxx_{output_id}.bin"
try:
ktool.cxx_index(merged_kzip, bin_file)
logger.info(
f"Successfully indexed program {args.package_name} to binary: {bin_file}"
)
except subprocess.CalledProcessError:
# TODO(Evan): For now, if this errors just keep going
logger.error(
f"Failed to index program {args.package_name} to binary: {bin_file}"
)
# Store the program into a graphml file
graphml = Path(td) / f"kythe_output_graphml_{output_id}.xml"
with open(graphml, "w") as fw, open(bin_file, "rb") as fr:
gs = GraphStorage()
gs.process_stream(fr)
fw.write(gs.to_graphml())
logger.info(
f"Successfully stored program {args.package_name} in graphml file: {graphml}"
)
# Copy graphml file
graphml_final = Path(td) / "graph.xml"
shutil.copy2(graphml, graphml_final)
logger.debug("Loading graphml file into JanusGraph...")
# TODO(Evan): This needs to wait until JanusGraph is ready. For some reason, even if the container is running and healthy, it's not ready to accept connections.
# Load graphml file into JanusGraph
from gremlin_python.process.anonymous_traversal import traversal
from gremlin_python.driver.driver_remote_connection import (
DriverRemoteConnection,
)
g = traversal().withRemote(
DriverRemoteConnection("ws://graphdb:8182/gremlin", "g")
)
g.io(str(graphml_final)).read().iterate()
logger.debug("Successfully loaded graphml file into JanusGraph")
return True
def serve(self):
"""Main loop to process tasks from queue"""
if self.task_queue is None:
raise ValueError("Task queue is not initialized")
if self.output_queue is None:
raise ValueError("Output queue is not initialized")
logger.info("Starting downloader service")
while True:
rq_item: Optional[RQItem] = self.task_queue.pop()
if rq_item is not None:
task_index: IndexRequest = rq_item.deserialized
success = self.process_task(task_index)
if success:
self.output_queue.push(
IndexOutput(
package_name=task_index.package_name,
sanitizer=task_index.sanitizer,
ossfuzz=task_index.ossfuzz,
source_path=task_index.source_path,
task_id=task_index.task_id,
build_type=task_index.build_type,
)
)
self.task_queue.ack_item(rq_item.item_id)
logger.info(f"Successfully processed task {task_index.task_id}")
else:
logger.error(f"Failed to process task {task_index.task_id}")
logger.info(f"Sleeping for {self.sleep_time} seconds")
time.sleep(self.sleep_time)
program-model/src/buttercup/program_model/settings.py
+63
View File
@@ -0,0 +1,63 @@
from pydantic_settings import BaseSettings, CliSubCommand
from pydantic import BaseModel, Field
from typing import Annotated
class Config:
cli_parse_args = True
nested_model_default_partial_update = True
env_nested_delimiter = "__"
extra = "allow"
env_prefix = "BUTTERCUP_PROGRAM_MODEL_"
class ButtercupBaseSettings(BaseSettings):
class Config:
cli_parse_args = True
nested_model_default_partial_update = True
env_nested_delimiter = "__"
extra = "allow"
class BuilderSettings(BaseModel):
allow_pull: Annotated[bool, Field(default=True, description="Allow pull")]
base_image_url: Annotated[
str, Field(default="gcr.io/oss-fuzz", description="Base image URL")
]
class WorkerSettings(BaseModel):
redis_url: Annotated[
str, Field(default="redis://127.0.0.1:6379", description="Redis URL")
]
sleep_time: Annotated[
float, Field(default=1.0, description="Sleep time between checks in seconds")
]
wdir: Annotated[str, Field(default=..., description="Working directory")]
python: Annotated[str, Field(default="python", description="Python path")]
class IndexerSettings(BaseModel):
kythe_dir: Annotated[
str, Field(default="scripts/gzs/kythe", description="Kythe directory")
]
script_dir: Annotated[str, Field(default="scripts", description="Script directory")]
class ProgramModelServeCommand(WorkerSettings, IndexerSettings, BuilderSettings):
pass
class ProgramModelProcessCommand(WorkerSettings, IndexerSettings, BuilderSettings):
package_name: Annotated[str, Field(description="Package name", default=...)]
sanitizer: Annotated[str, Field(description="Sanitizer", default=...)]
ossfuzz: Annotated[str, Field(description="OSSFuzz directory", default=...)]
source_path: Annotated[str, Field(description="Source path", default=...)]
task_id: Annotated[str, Field(description="Task ID", default=...)]
build_type: Annotated[str, Field(description="Build type", default=...)]
class ProgramModelSettings(ButtercupBaseSettings):
log_level: Annotated[str, Field(default="info", description="Log level")]
serve: CliSubCommand[ProgramModelServeCommand]
process: CliSubCommand[ProgramModelProcessCommand]
program-model/uv.lock
Generated
+1340
View File
File diff suppressed because it is too large Load Diff