sendnrw/netbird
2
0
Fork 0
mirror of https://github.com/netbirdio/netbird.git synced 2026-10-07 22:19:08 +02:00
Code Issues Packages Projects Releases Wiki Activity

merge main

Browse Source
This commit is contained in:
pascal
2026-07-28 14:14:07 +02:00
parent 90c45a6e24 1816a020c4
commit 0c2a6a8628
2033 changed files with 356952 additions and 30390 deletions
Download Patch File Download Diff File Expand all files Collapse all files
proxy/Dockerfile
+20
View File
@@ -0,0 +1,20 @@
FROM golang:1.25-alpine AS builder
WORKDIR /app
RUN echo "netbird:x:1000:1000:netbird:/var/lib/netbird:/sbin/nologin" > /tmp/passwd && \
echo "netbird:x:1000:netbird" > /tmp/group && \
mkdir -p /tmp/var/lib/netbird && \
mkdir -p /tmp/certs
FROM gcr.io/distroless/base:debug
ARG TARGETPLATFORM
COPY ${TARGETPLATFORM}/netbird-proxy /go/bin/netbird-proxy
COPY --from=builder /tmp/passwd /etc/passwd
COPY --from=builder /tmp/group /etc/group
COPY --from=builder --chown=1000:1000 /tmp/var/lib/netbird /var/lib/netbird
COPY --from=builder --chown=1000:1000 --chmod=755 /tmp/certs /certs
USER netbird:netbird
ENV HOME=/var/lib/netbird
ENV NB_PROXY_ADDRESS=":8443"
EXPOSE 8443
ENTRYPOINT ["/go/bin/netbird-proxy"]
proxy/Dockerfile.multistage
+39
View File
@@ -0,0 +1,39 @@
FROM golang:1.25-alpine AS builder
WORKDIR /app
COPY go.mod go.sum ./
RUN go mod download
COPY client ./client
COPY dns ./dns
COPY encryption ./encryption
COPY flow ./flow
COPY formatter ./formatter
COPY monotime ./monotime
COPY management ./management
COPY proxy ./proxy
COPY route ./route
COPY shared ./shared
COPY sharedsock ./sharedsock
COPY trustedproxy ./trustedproxy
COPY upload-server ./upload-server
COPY util ./util
COPY version ./version
RUN CGO_ENABLED=0 GOOS=linux go build -ldflags="-s -w" -o netbird-proxy ./proxy/cmd/proxy
RUN echo "netbird:x:1000:1000:netbird:/var/lib/netbird:/sbin/nologin" > /tmp/passwd && \
echo "netbird:x:1000:netbird" > /tmp/group && \
mkdir -p /tmp/var/lib/netbird && \
mkdir -p /tmp/certs
FROM gcr.io/distroless/base:debug
COPY --from=builder /app/netbird-proxy /usr/bin/netbird-proxy
COPY --from=builder /tmp/passwd /etc/passwd
COPY --from=builder /tmp/group /etc/group
COPY --from=builder --chown=1000:1000 /tmp/var/lib/netbird /var/lib/netbird
COPY --from=builder --chown=1000:1000 --chmod=755 /tmp/certs /certs
USER netbird:netbird
ENV HOME=/var/lib/netbird
ENV NB_PROXY_ADDRESS=":8443"
EXPOSE 8443
ENTRYPOINT ["/usr/bin/netbird-proxy"]
proxy/LICENSE
+661
View File
@@ -0,0 +1,661 @@
GNU AFFERO GENERAL PUBLIC LICENSE
Version 3, 19 November 2007
Copyright (C) 2007 Free Software Foundation, Inc. <https://fsf.org/>
Everyone is permitted to copy and distribute verbatim copies
of this license document, but changing it is not allowed.
Preamble
The GNU Affero General Public License is a free, copyleft license for
software and other kinds of works, specifically designed to ensure
cooperation with the community in the case of network server software.
The licenses for most software and other practical works are designed
to take away your freedom to share and change the works. By contrast,
our General Public Licenses are intended to guarantee your freedom to
share and change all versions of a program--to make sure it remains free
software for all its users.
When we speak of free software, we are referring to freedom, not
price. Our General Public Licenses are designed to make sure that you
have the freedom to distribute copies of free software (and charge for
them if you wish), that you receive source code or can get it if you
want it, that you can change the software or use pieces of it in new
free programs, and that you know you can do these things.
Developers that use our General Public Licenses protect your rights
with two steps: (1) assert copyright on the software, and (2) offer
you this License which gives you legal permission to copy, distribute
and/or modify the software.
A secondary benefit of defending all users' freedom is that
improvements made in alternate versions of the program, if they
receive widespread use, become available for other developers to
incorporate. Many developers of free software are heartened and
encouraged by the resulting cooperation. However, in the case of
software used on network servers, this result may fail to come about.
The GNU General Public License permits making a modified version and
letting the public access it on a server without ever releasing its
source code to the public.
The GNU Affero General Public License is designed specifically to
ensure that, in such cases, the modified source code becomes available
to the community. It requires the operator of a network server to
provide the source code of the modified version running there to the
users of that server. Therefore, public use of a modified version, on
a publicly accessible server, gives the public access to the source
code of the modified version.
An older license, called the Affero General Public License and
published by Affero, was designed to accomplish similar goals. This is
a different license, not a version of the Affero GPL, but Affero has
released a new version of the Affero GPL which permits relicensing under
this license.
The precise terms and conditions for copying, distribution and
modification follow.
TERMS AND CONDITIONS
0. Definitions.
"This License" refers to version 3 of the GNU Affero General Public License.
"Copyright" also means copyright-like laws that apply to other kinds of
works, such as semiconductor masks.
"The Program" refers to any copyrightable work licensed under this
License. Each licensee is addressed as "you". "Licensees" and
"recipients" may be individuals or organizations.
To "modify" a work means to copy from or adapt all or part of the work
in a fashion requiring copyright permission, other than the making of an
exact copy. The resulting work is called a "modified version" of the
earlier work or a work "based on" the earlier work.
A "covered work" means either the unmodified Program or a work based
on the Program.
To "propagate" a work means to do anything with it that, without
permission, would make you directly or secondarily liable for
infringement under applicable copyright law, except executing it on a
computer or modifying a private copy. Propagation includes copying,
distribution (with or without modification), making available to the
public, and in some countries other activities as well.
To "convey" a work means any kind of propagation that enables other
parties to make or receive copies. Mere interaction with a user through
a computer network, with no transfer of a copy, is not conveying.
An interactive user interface displays "Appropriate Legal Notices"
to the extent that it includes a convenient and prominently visible
feature that (1) displays an appropriate copyright notice, and (2)
tells the user that there is no warranty for the work (except to the
extent that warranties are provided), that licensees may convey the
work under this License, and how to view a copy of this License. If
the interface presents a list of user commands or options, such as a
menu, a prominent item in the list meets this criterion.
1. Source Code.
The "source code" for a work means the preferred form of the work
for making modifications to it. "Object code" means any non-source
form of a work.
A "Standard Interface" means an interface that either is an official
standard defined by a recognized standards body, or, in the case of
interfaces specified for a particular programming language, one that
is widely used among developers working in that language.
The "System Libraries" of an executable work include anything, other
than the work as a whole, that (a) is included in the normal form of
packaging a Major Component, but which is not part of that Major
Component, and (b) serves only to enable use of the work with that
Major Component, or to implement a Standard Interface for which an
implementation is available to the public in source code form. A
"Major Component", in this context, means a major essential component
(kernel, window system, and so on) of the specific operating system
(if any) on which the executable work runs, or a compiler used to
produce the work, or an object code interpreter used to run it.
The "Corresponding Source" for a work in object code form means all
the source code needed to generate, install, and (for an executable
work) run the object code and to modify the work, including scripts to
control those activities. However, it does not include the work's
System Libraries, or general-purpose tools or generally available free
programs which are used unmodified in performing those activities but
which are not part of the work. For example, Corresponding Source
includes interface definition files associated with source files for
the work, and the source code for shared libraries and dynamically
linked subprograms that the work is specifically designed to require,
such as by intimate data communication or control flow between those
subprograms and other parts of the work.
The Corresponding Source need not include anything that users
can regenerate automatically from other parts of the Corresponding
Source.
The Corresponding Source for a work in source code form is that
same work.
2. Basic Permissions.
All rights granted under this License are granted for the term of
copyright on the Program, and are irrevocable provided the stated
conditions are met. This License explicitly affirms your unlimited
permission to run the unmodified Program. The output from running a
covered work is covered by this License only if the output, given its
content, constitutes a covered work. This License acknowledges your
rights of fair use or other equivalent, as provided by copyright law.
You may make, run and propagate covered works that you do not
convey, without conditions so long as your license otherwise remains
in force. You may convey covered works to others for the sole purpose
of having them make modifications exclusively for you, or provide you
with facilities for running those works, provided that you comply with
the terms of this License in conveying all material for which you do
not control copyright. Those thus making or running the covered works
for you must do so exclusively on your behalf, under your direction
and control, on terms that prohibit them from making any copies of
your copyrighted material outside their relationship with you.
Conveying under any other circumstances is permitted solely under
the conditions stated below. Sublicensing is not allowed; section 10
makes it unnecessary.
3. Protecting Users' Legal Rights From Anti-Circumvention Law.
No covered work shall be deemed part of an effective technological
measure under any applicable law fulfilling obligations under article
11 of the WIPO copyright treaty adopted on 20 December 1996, or
similar laws prohibiting or restricting circumvention of such
measures.
When you convey a covered work, you waive any legal power to forbid
circumvention of technological measures to the extent such circumvention
is effected by exercising rights under this License with respect to
the covered work, and you disclaim any intention to limit operation or
modification of the work as a means of enforcing, against the work's
users, your or third parties' legal rights to forbid circumvention of
technological measures.
4. Conveying Verbatim Copies.
You may convey verbatim copies of the Program's source code as you
receive it, in any medium, provided that you conspicuously and
appropriately publish on each copy an appropriate copyright notice;
keep intact all notices stating that this License and any
non-permissive terms added in accord with section 7 apply to the code;
keep intact all notices of the absence of any warranty; and give all
recipients a copy of this License along with the Program.
You may charge any price or no price for each copy that you convey,
and you may offer support or warranty protection for a fee.
5. Conveying Modified Source Versions.
You may convey a work based on the Program, or the modifications to
produce it from the Program, in the form of source code under the
terms of section 4, provided that you also meet all of these conditions:
a) The work must carry prominent notices stating that you modified
it, and giving a relevant date.
b) The work must carry prominent notices stating that it is
released under this License and any conditions added under section
7. This requirement modifies the requirement in section 4 to
"keep intact all notices".
c) You must license the entire work, as a whole, under this
License to anyone who comes into possession of a copy. This
License will therefore apply, along with any applicable section 7
additional terms, to the whole of the work, and all its parts,
regardless of how they are packaged. This License gives no
permission to license the work in any other way, but it does not
invalidate such permission if you have separately received it.
d) If the work has interactive user interfaces, each must display
Appropriate Legal Notices; however, if the Program has interactive
interfaces that do not display Appropriate Legal Notices, your
work need not make them do so.
A compilation of a covered work with other separate and independent
works, which are not by their nature extensions of the covered work,
and which are not combined with it such as to form a larger program,
in or on a volume of a storage or distribution medium, is called an
"aggregate" if the compilation and its resulting copyright are not
used to limit the access or legal rights of the compilation's users
beyond what the individual works permit. Inclusion of a covered work
in an aggregate does not cause this License to apply to the other
parts of the aggregate.
6. Conveying Non-Source Forms.
You may convey a covered work in object code form under the terms
of sections 4 and 5, provided that you also convey the
machine-readable Corresponding Source under the terms of this License,
in one of these ways:
a) Convey the object code in, or embodied in, a physical product
(including a physical distribution medium), accompanied by the
Corresponding Source fixed on a durable physical medium
customarily used for software interchange.
b) Convey the object code in, or embodied in, a physical product
(including a physical distribution medium), accompanied by a
written offer, valid for at least three years and valid for as
long as you offer spare parts or customer support for that product
model, to give anyone who possesses the object code either (1) a
copy of the Corresponding Source for all the software in the
product that is covered by this License, on a durable physical
medium customarily used for software interchange, for a price no
more than your reasonable cost of physically performing this
conveying of source, or (2) access to copy the
Corresponding Source from a network server at no charge.
c) Convey individual copies of the object code with a copy of the
written offer to provide the Corresponding Source. This
alternative is allowed only occasionally and noncommercially, and
only if you received the object code with such an offer, in accord
with subsection 6b.
d) Convey the object code by offering access from a designated
place (gratis or for a charge), and offer equivalent access to the
Corresponding Source in the same way through the same place at no
further charge. You need not require recipients to copy the
Corresponding Source along with the object code. If the place to
copy the object code is a network server, the Corresponding Source
may be on a different server (operated by you or a third party)
that supports equivalent copying facilities, provided you maintain
clear directions next to the object code saying where to find the
Corresponding Source. Regardless of what server hosts the
Corresponding Source, you remain obligated to ensure that it is
available for as long as needed to satisfy these requirements.
e) Convey the object code using peer-to-peer transmission, provided
you inform other peers where the object code and Corresponding
Source of the work are being offered to the general public at no
charge under subsection 6d.
A separable portion of the object code, whose source code is excluded
from the Corresponding Source as a System Library, need not be
included in conveying the object code work.
A "User Product" is either (1) a "consumer product", which means any
tangible personal property which is normally used for personal, family,
or household purposes, or (2) anything designed or sold for incorporation
into a dwelling. In determining whether a product is a consumer product,
doubtful cases shall be resolved in favor of coverage. For a particular
product received by a particular user, "normally used" refers to a
typical or common use of that class of product, regardless of the status
of the particular user or of the way in which the particular user
actually uses, or expects or is expected to use, the product. A product
is a consumer product regardless of whether the product has substantial
commercial, industrial or non-consumer uses, unless such uses represent
the only significant mode of use of the product.
"Installation Information" for a User Product means any methods,
procedures, authorization keys, or other information required to install
and execute modified versions of a covered work in that User Product from
a modified version of its Corresponding Source. The information must
suffice to ensure that the continued functioning of the modified object
code is in no case prevented or interfered with solely because
modification has been made.
If you convey an object code work under this section in, or with, or
specifically for use in, a User Product, and the conveying occurs as
part of a transaction in which the right of possession and use of the
User Product is transferred to the recipient in perpetuity or for a
fixed term (regardless of how the transaction is characterized), the
Corresponding Source conveyed under this section must be accompanied
by the Installation Information. But this requirement does not apply
if neither you nor any third party retains the ability to install
modified object code on the User Product (for example, the work has
been installed in ROM).
The requirement to provide Installation Information does not include a
requirement to continue to provide support service, warranty, or updates
for a work that has been modified or installed by the recipient, or for
the User Product in which it has been modified or installed. Access to a
network may be denied when the modification itself materially and
adversely affects the operation of the network or violates the rules and
protocols for communication across the network.
Corresponding Source conveyed, and Installation Information provided,
in accord with this section must be in a format that is publicly
documented (and with an implementation available to the public in
source code form), and must require no special password or key for
unpacking, reading or copying.
7. Additional Terms.
"Additional permissions" are terms that supplement the terms of this
License by making exceptions from one or more of its conditions.
Additional permissions that are applicable to the entire Program shall
be treated as though they were included in this License, to the extent
that they are valid under applicable law. If additional permissions
apply only to part of the Program, that part may be used separately
under those permissions, but the entire Program remains governed by
this License without regard to the additional permissions.
When you convey a copy of a covered work, you may at your option
remove any additional permissions from that copy, or from any part of
it. (Additional permissions may be written to require their own
removal in certain cases when you modify the work.) You may place
additional permissions on material, added by you to a covered work,
for which you have or can give appropriate copyright permission.
Notwithstanding any other provision of this License, for material you
add to a covered work, you may (if authorized by the copyright holders of
that material) supplement the terms of this License with terms:
a) Disclaiming warranty or limiting liability differently from the
terms of sections 15 and 16 of this License; or
b) Requiring preservation of specified reasonable legal notices or
author attributions in that material or in the Appropriate Legal
Notices displayed by works containing it; or
c) Prohibiting misrepresentation of the origin of that material, or
requiring that modified versions of such material be marked in
reasonable ways as different from the original version; or
d) Limiting the use for publicity purposes of names of licensors or
authors of the material; or
e) Declining to grant rights under trademark law for use of some
trade names, trademarks, or service marks; or
f) Requiring indemnification of licensors and authors of that
material by anyone who conveys the material (or modified versions of
it) with contractual assumptions of liability to the recipient, for
any liability that these contractual assumptions directly impose on
those licensors and authors.
All other non-permissive additional terms are considered "further
restrictions" within the meaning of section 10. If the Program as you
received it, or any part of it, contains a notice stating that it is
governed by this License along with a term that is a further
restriction, you may remove that term. If a license document contains
a further restriction but permits relicensing or conveying under this
License, you may add to a covered work material governed by the terms
of that license document, provided that the further restriction does
not survive such relicensing or conveying.
If you add terms to a covered work in accord with this section, you
must place, in the relevant source files, a statement of the
additional terms that apply to those files, or a notice indicating
where to find the applicable terms.
Additional terms, permissive or non-permissive, may be stated in the
form of a separately written license, or stated as exceptions;
the above requirements apply either way.
8. Termination.
You may not propagate or modify a covered work except as expressly
provided under this License. Any attempt otherwise to propagate or
modify it is void, and will automatically terminate your rights under
this License (including any patent licenses granted under the third
paragraph of section 11).
However, if you cease all violation of this License, then your
license from a particular copyright holder is reinstated (a)
provisionally, unless and until the copyright holder explicitly and
finally terminates your license, and (b) permanently, if the copyright
holder fails to notify you of the violation by some reasonable means
prior to 60 days after the cessation.
Moreover, your license from a particular copyright holder is
reinstated permanently if the copyright holder notifies you of the
violation by some reasonable means, this is the first time you have
received notice of violation of this License (for any work) from that
copyright holder, and you cure the violation prior to 30 days after
your receipt of the notice.
Termination of your rights under this section does not terminate the
licenses of parties who have received copies or rights from you under
this License. If your rights have been terminated and not permanently
reinstated, you do not qualify to receive new licenses for the same
material under section 10.
9. Acceptance Not Required for Having Copies.
You are not required to accept this License in order to receive or
run a copy of the Program. Ancillary propagation of a covered work
occurring solely as a consequence of using peer-to-peer transmission
to receive a copy likewise does not require acceptance. However,
nothing other than this License grants you permission to propagate or
modify any covered work. These actions infringe copyright if you do
not accept this License. Therefore, by modifying or propagating a
covered work, you indicate your acceptance of this License to do so.
10. Automatic Licensing of Downstream Recipients.
Each time you convey a covered work, the recipient automatically
receives a license from the original licensors, to run, modify and
propagate that work, subject to this License. You are not responsible
for enforcing compliance by third parties with this License.
An "entity transaction" is a transaction transferring control of an
organization, or substantially all assets of one, or subdividing an
organization, or merging organizations. If propagation of a covered
work results from an entity transaction, each party to that
transaction who receives a copy of the work also receives whatever
licenses to the work the party's predecessor in interest had or could
give under the previous paragraph, plus a right to possession of the
Corresponding Source of the work from the predecessor in interest, if
the predecessor has it or can get it with reasonable efforts.
You may not impose any further restrictions on the exercise of the
rights granted or affirmed under this License. For example, you may
not impose a license fee, royalty, or other charge for exercise of
rights granted under this License, and you may not initiate litigation
(including a cross-claim or counterclaim in a lawsuit) alleging that
any patent claim is infringed by making, using, selling, offering for
sale, or importing the Program or any portion of it.
11. Patents.
A "contributor" is a copyright holder who authorizes use under this
License of the Program or a work on which the Program is based. The
work thus licensed is called the contributor's "contributor version".
A contributor's "essential patent claims" are all patent claims
owned or controlled by the contributor, whether already acquired or
hereafter acquired, that would be infringed by some manner, permitted
by this License, of making, using, or selling its contributor version,
but do not include claims that would be infringed only as a
consequence of further modification of the contributor version. For
purposes of this definition, "control" includes the right to grant
patent sublicenses in a manner consistent with the requirements of
this License.
Each contributor grants you a non-exclusive, worldwide, royalty-free
patent license under the contributor's essential patent claims, to
make, use, sell, offer for sale, import and otherwise run, modify and
propagate the contents of its contributor version.
In the following three paragraphs, a "patent license" is any express
agreement or commitment, however denominated, not to enforce a patent
(such as an express permission to practice a patent or covenant not to
sue for patent infringement). To "grant" such a patent license to a
party means to make such an agreement or commitment not to enforce a
patent against the party.
If you convey a covered work, knowingly relying on a patent license,
and the Corresponding Source of the work is not available for anyone
to copy, free of charge and under the terms of this License, through a
publicly available network server or other readily accessible means,
then you must either (1) cause the Corresponding Source to be so
available, or (2) arrange to deprive yourself of the benefit of the
patent license for this particular work, or (3) arrange, in a manner
consistent with the requirements of this License, to extend the patent
license to downstream recipients. "Knowingly relying" means you have
actual knowledge that, but for the patent license, your conveying the
covered work in a country, or your recipient's use of the covered work
in a country, would infringe one or more identifiable patents in that
country that you have reason to believe are valid.
If, pursuant to or in connection with a single transaction or
arrangement, you convey, or propagate by procuring conveyance of, a
covered work, and grant a patent license to some of the parties
receiving the covered work authorizing them to use, propagate, modify
or convey a specific copy of the covered work, then the patent license
you grant is automatically extended to all recipients of the covered
work and works based on it.
A patent license is "discriminatory" if it does not include within
the scope of its coverage, prohibits the exercise of, or is
conditioned on the non-exercise of one or more of the rights that are
specifically granted under this License. You may not convey a covered
work if you are a party to an arrangement with a third party that is
in the business of distributing software, under which you make payment
to the third party based on the extent of your activity of conveying
the work, and under which the third party grants, to any of the
parties who would receive the covered work from you, a discriminatory
patent license (a) in connection with copies of the covered work
conveyed by you (or copies made from those copies), or (b) primarily
for and in connection with specific products or compilations that
contain the covered work, unless you entered into that arrangement,
or that patent license was granted, prior to 28 March 2007.
Nothing in this License shall be construed as excluding or limiting
any implied license or other defenses to infringement that may
otherwise be available to you under applicable patent law.
12. No Surrender of Others' Freedom.
If conditions are imposed on you (whether by court order, agreement or
otherwise) that contradict the conditions of this License, they do not
excuse you from the conditions of this License. If you cannot convey a
covered work so as to satisfy simultaneously your obligations under this
License and any other pertinent obligations, then as a consequence you may
not convey it at all. For example, if you agree to terms that obligate you
to collect a royalty for further conveying from those to whom you convey
the Program, the only way you could satisfy both those terms and this
License would be to refrain entirely from conveying the Program.
13. Remote Network Interaction; Use with the GNU General Public License.
Notwithstanding any other provision of this License, if you modify the
Program, your modified version must prominently offer all users
interacting with it remotely through a computer network (if your version
supports such interaction) an opportunity to receive the Corresponding
Source of your version by providing access to the Corresponding Source
from a network server at no charge, through some standard or customary
means of facilitating copying of software. This Corresponding Source
shall include the Corresponding Source for any work covered by version 3
of the GNU General Public License that is incorporated pursuant to the
following paragraph.
Notwithstanding any other provision of this License, you have
permission to link or combine any covered work with a work licensed
under version 3 of the GNU General Public License into a single
combined work, and to convey the resulting work. The terms of this
License will continue to apply to the part which is the covered work,
but the work with which it is combined will remain governed by version
3 of the GNU General Public License.
14. Revised Versions of this License.
The Free Software Foundation may publish revised and/or new versions of
the GNU Affero General Public License from time to time. Such new versions
will be similar in spirit to the present version, but may differ in detail to
address new problems or concerns.
Each version is given a distinguishing version number. If the
Program specifies that a certain numbered version of the GNU Affero General
Public License "or any later version" applies to it, you have the
option of following the terms and conditions either of that numbered
version or of any later version published by the Free Software
Foundation. If the Program does not specify a version number of the
GNU Affero General Public License, you may choose any version ever published
by the Free Software Foundation.
If the Program specifies that a proxy can decide which future
versions of the GNU Affero General Public License can be used, that proxy's
public statement of acceptance of a version permanently authorizes you
to choose that version for the Program.
Later license versions may give you additional or different
permissions. However, no additional obligations are imposed on any
author or copyright holder as a result of your choosing to follow a
later version.
15. Disclaimer of Warranty.
THERE IS NO WARRANTY FOR THE PROGRAM, TO THE EXTENT PERMITTED BY
APPLICABLE LAW. EXCEPT WHEN OTHERWISE STATED IN WRITING THE COPYRIGHT
HOLDERS AND/OR OTHER PARTIES PROVIDE THE PROGRAM "AS IS" WITHOUT WARRANTY
OF ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING, BUT NOT LIMITED TO,
THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
PURPOSE. THE ENTIRE RISK AS TO THE QUALITY AND PERFORMANCE OF THE PROGRAM
IS WITH YOU. SHOULD THE PROGRAM PROVE DEFECTIVE, YOU ASSUME THE COST OF
ALL NECESSARY SERVICING, REPAIR OR CORRECTION.
16. Limitation of Liability.
IN NO EVENT UNLESS REQUIRED BY APPLICABLE LAW OR AGREED TO IN WRITING
WILL ANY COPYRIGHT HOLDER, OR ANY OTHER PARTY WHO MODIFIES AND/OR CONVEYS
THE PROGRAM AS PERMITTED ABOVE, BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY
GENERAL, SPECIAL, INCIDENTAL OR CONSEQUENTIAL DAMAGES ARISING OUT OF THE
USE OR INABILITY TO USE THE PROGRAM (INCLUDING BUT NOT LIMITED TO LOSS OF
DATA OR DATA BEING RENDERED INACCURATE OR LOSSES SUSTAINED BY YOU OR THIRD
PARTIES OR A FAILURE OF THE PROGRAM TO OPERATE WITH ANY OTHER PROGRAMS),
EVEN IF SUCH HOLDER OR OTHER PARTY HAS BEEN ADVISED OF THE POSSIBILITY OF
SUCH DAMAGES.
17. Interpretation of Sections 15 and 16.
If the disclaimer of warranty and limitation of liability provided
above cannot be given local legal effect according to their terms,
reviewing courts shall apply local law that most closely approximates
an absolute waiver of all civil liability in connection with the
Program, unless a warranty or assumption of liability accompanies a
copy of the Program in return for a fee.
END OF TERMS AND CONDITIONS
How to Apply These Terms to Your New Programs
If you develop a new program, and you want it to be of the greatest
possible use to the public, the best way to achieve this is to make it
free software which everyone can redistribute and change under these terms.
To do so, attach the following notices to the program. It is safest
to attach them to the start of each source file to most effectively
state the exclusion of warranty; and each file should have at least
the "copyright" line and a pointer to where the full notice is found.
<one line to give the program's name and a brief idea of what it does.>
Copyright (C) <year> <name of author>
This program is free software: you can redistribute it and/or modify
it under the terms of the GNU Affero General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.
This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU Affero General Public License for more details.
You should have received a copy of the GNU Affero General Public License
along with this program. If not, see <https://www.gnu.org/licenses/>.
Also add information on how to contact you by electronic and paper mail.
If your software can interact with users remotely through a computer
network, you should also make sure that it provides a way for users to
get its source. For example, if your program is a web application, its
interface could display a "Source" link that leads users to an archive
of the code. There are many ways you could offer source, and different
solutions will be better for different programs; see section 13 for the
specific requirements.
You should also get your employer (if you work as a programmer) or school,
if any, to sign a "copyright disclaimer" for the program, if necessary.
For more information on this, and how to apply and follow the GNU AGPL, see
<https://www.gnu.org/licenses/>.
proxy/README.md
+80
View File
@@ -0,0 +1,80 @@
# Netbird Reverse Proxy
The NetBird Reverse Proxy is a separate service that can act as a public entrypoint to certain resources within a NetBird network.
At a high level, the way that it operates is:
- Configured routes are communicated from the Management server to the proxy.
- For each route the proxy creates a NetBird connection to the NetBird Peer that hosts the resource.
- When traffic hits the proxy at the address and path configured for the proxied resource, the NetBird Proxy brings up a relevant authentication method for that resource.
- On successful authentication the proxy will forward traffic onwards to the NetBird Peer.
Proxy Authentication methods supported are:
- No authentication
- Oauth2/OIDC
- Emailed Magic Link
- Simple PIN
- HTTP Basic Auth Username and Password
## Management Connection and Authentication
The Proxy communicates with the Management server over a gRPC connection.
Proxies act as clients to the Management server, the following RPCs are used:
- Server-side streaming for proxied service updates.
- Client-side streaming for proxy logs.
To authenticate with the Management server, the proxy server uses Machine-to-Machine OAuth2.
If you are using the embedded IdP //TODO: explain how to get credentials.
Otherwise, create a new machine-to-machine profile in your IdP for proxy servers and set the relevant settings in the proxy's environment or flags (see below).
## User Authentication
When a request hits the Proxy, it looks up the permitted authentication methods for the Host domain.
If no authentication methods are registered for the Host domain, then no authentication will be applied (for fully public resources).
If any authentication methods are registered for the Host domain, then the Proxy will first serve an authentication page allowing the user to select an authentication method (from the permitted methods) and enter the required information for that authentication method.
If the user is successfully authenticated, their request will be forwarded through to the Proxy to be proxied to the relevant Peer.
Successful authentication does not guarantee a successful forwarding of the request as there may be failures behind the Proxy, such as with Peer connectivity or the underlying resource.
## TLS
Due to the authentication provided, the Proxy uses HTTPS for its endpoint, even if the underlying service is HTTP.
Certificate generation can either be via ACME (by default, using Let's Encrypt, but alternative ACME providers can be used) or through certificate files.
When not using ACME, the proxy server attempts to load a certificate and key from the files `tls.crt` and `tls.key` in a specified certificate directory.
When using ACME, the proxy server will store generated certificates in the specified certificate directory.
## Auth UI
The authentication UI is a Vite + React application located in the `web/` directory. It is embedded into the Go binary at build time.
To build the UI:
```bash
cd web
npm install
npm run build
```
For UI development with hot reload (served at http://localhost:3031):
```bash
npm run dev
```
The built assets in `web/dist/` are embedded via `//go:embed` and served by the `web.ServeHTTP` handler.
## Configuration
NetBird Proxy deployment configuration is via flags or environment variables, with flags taking precedence over the environment.
The following deployment configuration is available:
| Flag | Env | Purpose | Default |
|------------------|----------------------------------|------------------------------------------------------------------------------------------------------------------------------------|----------------------------------------------------|
| `-debug` | `NB_PROXY_DEBUG_LOGS` | Enable debug logging | `false` |
| `-mgmt` | `NB_PROXY_MANAGEMENT_ADDRESS` | The address of the management server for the proxy to get configuration from. | `"https://api.netbird.io:443"` |
| `-addr` | `NB_PROXY_ADDRESS` | The address that the reverse proxy will listen on. | `":443` |
| `-url` | `NB_PROXY_URL` | The URL that the proxy will be reached at (where endpoints will be CNAMEd to). If unset, this will fall back to the proxy address. | `"proxy.netbird.io"` |
| `-cert-dir` | `NB_PROXY_CERTIFICATE_DIRECTORY` | The location that certificates are stored in. | `"./certs"` |
| `-acme-certs` | `NB_PROXY_ACME_CERTIFICATES` | Whether to use ACME to generate certificates. | `false` |
| `-acme-addr` | `NB_PROXY_ACME_ADDRESS` | The HTTP address the proxy will listen on to respond to HTTP-01 ACME challenges | `":80"` |
| `-acme-dir` | `NB_PROXY_ACME_DIRECTORY` | The directory URL of the ACME server to be used | `"https://acme-v02.api.letsencrypt.org/directory"` |
| `-oidc-id` | `NB_PROXY_OIDC_CLIENT_ID` | The OAuth2 Client ID for OIDC User Authentication | `"netbird-proxy"` |
| `-oidc-secret` | `NB_PROXY_OIDC_CLIENT_SECRET` | The OAuth2 Client Secret for OIDC User Authentication | `""` |
| `-oidc-endpoint` | `NB_PROXY_OIDC_ENDPOINT` | The OAuth2 provider endpoint for OIDC User Authentication | `"https://api.netbird.io/oauth2"` |
| `-oidc-scopes` | `NB_PROXY_OIDC_SCOPES` | The OAuth2 scopes for OIDC User Authentication, comma separated | `"openid,profile,email"` |
proxy/auth/auth.go
+107
View File
@@ -0,0 +1,107 @@
// Package auth contains exported proxy auth values.
// These are used to ensure coherent usage across management and proxy implementations.
package auth
import (
"crypto/ed25519"
"crypto/tls"
"fmt"
"time"
"github.com/golang-jwt/jwt/v5"
)
type Method string
const (
MethodPassword Method = "password"
MethodPIN Method = "pin"
MethodOIDC Method = "oidc"
MethodHeader Method = "header"
)
func (m Method) String() string {
return string(m)
}
const (
SessionCookieName = "nb_session"
DefaultSessionExpiry = 24 * time.Hour
SessionJWTIssuer = "netbird-management"
)
// ResolveProto determines the protocol scheme based on the forwarded proto
// configuration. When set to "http" or "https" the value is used directly.
// Otherwise TLS state is used: if conn is non-nil "https" is returned, else "http".
func ResolveProto(forwardedProto string, conn *tls.ConnectionState) string {
switch forwardedProto {
case "http", "https":
return forwardedProto
default:
if conn != nil {
return "https"
}
return "http"
}
}
// ValidateSessionJWT validates a session JWT and returns the user ID, the
// user's email (when carried), the authentication method, any embedded
// group memberships, and the parallel group display names. email,
// groups, and groupNames may be empty for tokens minted before those
// claims were introduced. groupNames pairs positionally with groups.
func ValidateSessionJWT(tokenString, domain string, publicKey ed25519.PublicKey) (userID, email, method string, groups, groupNames []string, err error) {
if publicKey == nil {
return "", "", "", nil, nil, fmt.Errorf("no public key configured for domain")
}
token, err := jwt.Parse(tokenString, func(t *jwt.Token) (interface{}, error) {
if _, ok := t.Method.(*jwt.SigningMethodEd25519); !ok {
return nil, fmt.Errorf("unexpected signing method: %v", t.Header["alg"])
}
return publicKey, nil
}, jwt.WithAudience(domain), jwt.WithIssuer(SessionJWTIssuer))
if err != nil {
return "", "", "", nil, nil, fmt.Errorf("parse token: %w", err)
}
claims, ok := token.Claims.(jwt.MapClaims)
if !ok || !token.Valid {
return "", "", "", nil, nil, fmt.Errorf("invalid token claims")
}
sub, _ := claims.GetSubject()
if sub == "" {
return "", "", "", nil, nil, fmt.Errorf("missing subject claim")
}
methodClaim, _ := claims["method"].(string)
emailClaim, _ := claims["email"].(string)
groups = extractGroupsClaim(claims["groups"])
groupNames = extractGroupsClaim(claims["group_names"])
return sub, emailClaim, methodClaim, groups, groupNames, nil
}
// extractGroupsClaim decodes the "groups" claim into a string slice. The JWT
// library decodes JSON arrays as []interface{}, so we coerce element-wise
// and skip non-string entries silently.
func extractGroupsClaim(claim interface{}) []string {
raw, ok := claim.([]interface{})
if !ok {
return nil
}
if len(raw) == 0 {
return nil
}
groups := make([]string, 0, len(raw))
for _, v := range raw {
if s, ok := v.(string); ok && s != "" {
groups = append(groups, s)
}
}
if len(groups) == 0 {
return nil
}
return groups
}
proxy/cmd/proxy/cmd/debug.go
+332
View File
@@ -0,0 +1,332 @@
package cmd
import (
"fmt"
"os"
"os/signal"
"path/filepath"
"strconv"
"strings"
"syscall"
"time"
"github.com/spf13/cobra"
"github.com/netbirdio/netbird/proxy/internal/debug"
)
var (
debugAddr string
jsonOutput bool
// status filters
statusFilterByIPs []string
statusFilterByNames []string
statusFilterByStatus string
statusFilterByConnectionType string
)
var debugCmd = &cobra.Command{
Use: "debug",
Short: "Debug commands for inspecting proxy state",
Long: "Debug commands for inspecting the reverse proxy state via the debug HTTP endpoint.",
}
var debugHealthCmd = &cobra.Command{
Use: "health",
Short: "Show proxy health status",
RunE: runDebugHealth,
SilenceUsage: true,
}
var debugClientsCmd = &cobra.Command{
Use: "clients",
Aliases: []string{"list"},
Short: "List all connected clients",
RunE: runDebugClients,
SilenceUsage: true,
}
var debugStatusCmd = &cobra.Command{
Use: "status <account-id>",
Short: "Show client status",
Args: cobra.ExactArgs(1),
RunE: runDebugStatus,
SilenceUsage: true,
}
var debugSyncCmd = &cobra.Command{
Use: "sync-response <account-id>",
Short: "Show client sync response",
Args: cobra.ExactArgs(1),
RunE: runDebugSync,
SilenceUsage: true,
}
var (
pingTimeout time.Duration
pingIPv4 bool
pingIPv6 bool
)
var debugPingCmd = &cobra.Command{
Use: "ping <account-id> <host> [port]",
Short: "TCP ping through a client",
Long: "Perform a TCP ping through a client's network to test connectivity.\nPort defaults to 80 if not specified.",
Args: cobra.RangeArgs(2, 3),
RunE: runDebugPing,
SilenceUsage: true,
}
var debugLogCmd = &cobra.Command{
Use: "log",
Short: "Manage client logging",
Long: "Commands to manage logging settings for a client connected through the proxy.",
}
var debugLogLevelCmd = &cobra.Command{
Use: "level <account-id> <level>",
Short: "Set client log level",
Long: "Set the log level for a client (trace, debug, info, warn, error).",
Args: cobra.ExactArgs(2),
RunE: runDebugLogLevel,
SilenceUsage: true,
}
var debugStartCmd = &cobra.Command{
Use: "start <account-id>",
Short: "Start a client",
Args: cobra.ExactArgs(1),
RunE: runDebugStart,
SilenceUsage: true,
}
var debugStopCmd = &cobra.Command{
Use: "stop <account-id>",
Short: "Stop a client",
Args: cobra.ExactArgs(1),
RunE: runDebugStop,
SilenceUsage: true,
}
var debugPerfCmd = &cobra.Command{
Use: "perf <pool-cap>",
Short: "Live-retune the tunnel buffer pool cap on all running clients",
Args: cobra.ExactArgs(1),
RunE: runDebugPerfSet,
SilenceUsage: true,
}
var debugRuntimeCmd = &cobra.Command{
Use: "runtime",
Short: "Show runtime stats (heap, goroutines, RSS)",
Args: cobra.NoArgs,
RunE: runDebugRuntime,
SilenceUsage: true,
}
var debugCaptureCmd = &cobra.Command{
Use: "capture <account-id> [filter expression]",
Short: "Capture packets on a client's WireGuard interface",
Long: `Captures decrypted packets flowing through a client's WireGuard interface.
Default output is human-readable text. Use --pcap or --output for pcap binary.
Filter arguments after the account ID use BPF-like syntax.
Examples:
netbird-proxy debug capture <account-id>
netbird-proxy debug capture <account-id> --duration 1m host 10.0.0.1
netbird-proxy debug capture <account-id> host 10.0.0.1 and tcp port 443
netbird-proxy debug capture <account-id> not port 22
netbird-proxy debug capture <account-id> -o capture.pcap
netbird-proxy debug capture <account-id> --pcap | tcpdump -r - -n
netbird-proxy debug capture <account-id> --pcap | tshark -r -`,
Args: cobra.MinimumNArgs(1),
RunE: runDebugCapture,
SilenceUsage: true,
}
func init() {
debugCmd.PersistentFlags().StringVar(&debugAddr, "addr", envStringOrDefault("NB_PROXY_DEBUG_ADDRESS", "localhost:8444"), "Debug endpoint address")
debugCmd.PersistentFlags().BoolVar(&jsonOutput, "json", false, "Output JSON instead of pretty format")
debugStatusCmd.Flags().StringSliceVar(&statusFilterByIPs, "filter-by-ips", nil, "Filter by peer IPs (comma-separated)")
debugStatusCmd.Flags().StringSliceVar(&statusFilterByNames, "filter-by-names", nil, "Filter by peer names (comma-separated)")
debugStatusCmd.Flags().StringVar(&statusFilterByStatus, "filter-by-status", "", "Filter by status (idle|connecting|connected)")
debugStatusCmd.Flags().StringVar(&statusFilterByConnectionType, "filter-by-connection-type", "", "Filter by connection type (P2P|Relayed)")
debugPingCmd.Flags().DurationVar(&pingTimeout, "timeout", 0, "Ping timeout (e.g., 10s)")
debugPingCmd.Flags().BoolVarP(&pingIPv4, "ipv4", "4", false, "Force IPv4")
debugPingCmd.Flags().BoolVarP(&pingIPv6, "ipv6", "6", false, "Force IPv6")
debugPingCmd.MarkFlagsMutuallyExclusive("ipv4", "ipv6")
debugCaptureCmd.Flags().DurationP("duration", "d", 0, "Capture duration (0 = server default)")
debugCaptureCmd.Flags().Bool("pcap", false, "Force pcap binary output (default when --output is set)")
debugCaptureCmd.Flags().BoolP("verbose", "v", false, "Show seq/ack, TTL, window, total length (text mode)")
debugCaptureCmd.Flags().Bool("ascii", false, "Print payload as ASCII after each packet (text mode)")
debugCaptureCmd.Flags().StringP("output", "o", "", "Write pcap to file instead of stdout")
debugCmd.AddCommand(debugHealthCmd)
debugCmd.AddCommand(debugClientsCmd)
debugCmd.AddCommand(debugStatusCmd)
debugCmd.AddCommand(debugSyncCmd)
debugCmd.AddCommand(debugPingCmd)
debugLogCmd.AddCommand(debugLogLevelCmd)
debugCmd.AddCommand(debugLogCmd)
debugCmd.AddCommand(debugStartCmd)
debugCmd.AddCommand(debugStopCmd)
debugCmd.AddCommand(debugPerfCmd)
debugCmd.AddCommand(debugRuntimeCmd)
debugCmd.AddCommand(debugCaptureCmd)
rootCmd.AddCommand(debugCmd)
}
func getDebugClient(cmd *cobra.Command) *debug.Client {
return debug.NewClient(debugAddr, jsonOutput, cmd.OutOrStdout())
}
func runDebugHealth(cmd *cobra.Command, _ []string) error {
return getDebugClient(cmd).Health(cmd.Context())
}
func runDebugClients(cmd *cobra.Command, _ []string) error {
return getDebugClient(cmd).ListClients(cmd.Context())
}
func runDebugStatus(cmd *cobra.Command, args []string) error {
return getDebugClient(cmd).ClientStatus(cmd.Context(), args[0], debug.StatusFilters{
IPs: statusFilterByIPs,
Names: statusFilterByNames,
Status: statusFilterByStatus,
ConnectionType: statusFilterByConnectionType,
})
}
func runDebugSync(cmd *cobra.Command, args []string) error {
return getDebugClient(cmd).ClientSyncResponse(cmd.Context(), args[0])
}
func runDebugPing(cmd *cobra.Command, args []string) error {
port := 80
if len(args) > 2 {
p, err := strconv.Atoi(args[2])
if err != nil {
return fmt.Errorf("invalid port: %w", err)
}
port = p
}
var ipVersion string
switch {
case pingIPv4:
ipVersion = "4"
case pingIPv6:
ipVersion = "6"
}
return getDebugClient(cmd).PingTCP(cmd.Context(), args[0], args[1], port, pingTimeout, ipVersion)
}
func runDebugLogLevel(cmd *cobra.Command, args []string) error {
return getDebugClient(cmd).SetLogLevel(cmd.Context(), args[0], args[1])
}
func runDebugStart(cmd *cobra.Command, args []string) error {
return getDebugClient(cmd).StartClient(cmd.Context(), args[0])
}
func runDebugStop(cmd *cobra.Command, args []string) error {
return getDebugClient(cmd).StopClient(cmd.Context(), args[0])
}
func runDebugPerfSet(cmd *cobra.Command, args []string) error {
n, err := strconv.ParseUint(args[0], 10, 32)
if err != nil {
return fmt.Errorf("invalid value %q: %w", args[0], err)
}
return getDebugClient(cmd).PerfSet(cmd.Context(), uint32(n))
}
func runDebugRuntime(cmd *cobra.Command, _ []string) error {
return getDebugClient(cmd).Runtime(cmd.Context())
}
func runDebugCapture(cmd *cobra.Command, args []string) error {
duration, _ := cmd.Flags().GetDuration("duration")
forcePcap, _ := cmd.Flags().GetBool("pcap")
verbose, _ := cmd.Flags().GetBool("verbose")
ascii, _ := cmd.Flags().GetBool("ascii")
outPath, _ := cmd.Flags().GetString("output")
// Default to text. Use pcap when --pcap is set or --output is given.
wantText := !forcePcap && outPath == ""
var filterExpr string
if len(args) > 1 {
filterExpr = strings.Join(args[1:], " ")
}
ctx, cancel := signal.NotifyContext(cmd.Context(), syscall.SIGINT, syscall.SIGTERM)
defer cancel()
out, cleanup, err := captureOutputWriter(cmd, outPath)
if err != nil {
return err
}
defer cleanup()
if wantText {
cmd.PrintErrln("Capturing packets... Press Ctrl+C to stop.")
} else {
cmd.PrintErrln("Capturing packets (pcap)... Press Ctrl+C to stop.")
}
var durationStr string
if duration > 0 {
durationStr = duration.String()
}
err = getDebugClient(cmd).Capture(ctx, debug.CaptureOptions{
AccountID: args[0],
Duration: durationStr,
FilterExpr: filterExpr,
Text: wantText,
Verbose: verbose,
ASCII: ascii,
Output: out,
})
if err != nil {
return err
}
cmd.PrintErrln("\nCapture finished.")
return nil
}
// captureOutputWriter returns the writer and cleanup function for capture output.
func captureOutputWriter(cmd *cobra.Command, outPath string) (out *os.File, cleanup func(), err error) {
if outPath != "" {
f, err := os.CreateTemp(filepath.Dir(outPath), filepath.Base(outPath)+".*.tmp")
if err != nil {
return nil, nil, fmt.Errorf("create output file: %w", err)
}
tmpPath := f.Name()
return f, func() {
if err := f.Close(); err != nil {
cmd.PrintErrf("close output file: %v\n", err)
}
if fi, err := os.Stat(tmpPath); err == nil && fi.Size() > 0 {
if err := os.Rename(tmpPath, outPath); err != nil {
cmd.PrintErrf("rename output file: %v\n", err)
} else {
cmd.PrintErrf("Wrote %s\n", outPath)
}
} else {
os.Remove(tmpPath)
}
}, nil
}
return os.Stdout, func() {
// no cleanup needed for stdout
}, nil
}
proxy/cmd/proxy/cmd/root.go
+307
View File
@@ -0,0 +1,307 @@
package cmd
import (
"context"
"fmt"
"os"
"os/signal"
"strconv"
"syscall"
"time"
log "github.com/sirupsen/logrus"
"github.com/spf13/cobra"
"golang.org/x/crypto/acme"
"github.com/netbirdio/netbird/shared/management/domain"
"github.com/netbirdio/netbird/client/embed"
"github.com/netbirdio/netbird/proxy"
nbacme "github.com/netbirdio/netbird/proxy/internal/acme"
"github.com/netbirdio/netbird/trustedproxy"
"github.com/netbirdio/netbird/util"
)
const (
// envPreallocatedBuffers caps the per-tunnel buffer pool. Zero (unset)
// keeps the upstream uncapped default.
envPreallocatedBuffers = "NB_PROXY_PREALLOCATED_BUFFERS"
// envMaxBatchSize overrides the per-tunnel batch size, which controls
// how many buffers each receive/TUN worker eagerly allocates. Zero
// (unset) keeps the platform default.
envMaxBatchSize = "NB_PROXY_MAX_BATCH_SIZE"
)
const DefaultManagementURL = "https://api.netbird.io:443"
// envProxyToken is the environment variable name for the proxy access token.
//
//nolint:gosec
const envProxyToken = "NB_PROXY_TOKEN"
var (
Version = "dev"
Commit = "unknown"
BuildDate = "unknown"
GoVersion = "unknown"
)
var (
logLevel string
debugLogs bool
mgmtAddr string
addr string
proxyDomain string
maxDialTimeout time.Duration
maxSessionIdleTimeout time.Duration
certDir string
acmeCerts bool
acmeAddr string
acmeDir string
acmeEABKID string
acmeEABHMACKey string
acmeChallengeType string
debugEndpoint bool
debugEndpointAddr string
healthAddr string
forwardedProto string
trustedProxies string
certFile string
certKeyFile string
certLockMethod string
wildcardCertDir string
wgPort uint16
proxyProtocol bool
preSharedKey string
supportsCustomPorts bool
requireSubdomain bool
private bool
geoDataDir string
crowdsecAPIURL string
crowdsecAPIKey string
)
var rootCmd = &cobra.Command{
Use: "proxy",
Short: "NetBird reverse proxy server",
Long: "NetBird reverse proxy server for proxying traffic to NetBird networks.",
Version: Version,
SilenceUsage: true,
RunE: runServer,
}
func init() {
rootCmd.PersistentFlags().StringVar(&logLevel, "log-level", envStringOrDefault("NB_PROXY_LOG_LEVEL", "info"), "Log level: panic, fatal, error, warn, info, debug, trace")
rootCmd.PersistentFlags().BoolVar(&debugLogs, "debug", envBoolOrDefault("NB_PROXY_DEBUG_LOGS", false), "Enable debug logs")
_ = rootCmd.PersistentFlags().MarkDeprecated("debug", "use --log-level instead")
rootCmd.Flags().StringVar(&mgmtAddr, "mgmt", envStringOrDefault("NB_PROXY_MANAGEMENT_ADDRESS", DefaultManagementURL), "Management address to connect to")
rootCmd.Flags().StringVar(&addr, "addr", envStringOrDefault("NB_PROXY_ADDRESS", ":443"), "Reverse proxy address to listen on")
rootCmd.Flags().StringVar(&proxyDomain, "domain", envStringOrDefault("NB_PROXY_DOMAIN", ""), "The Domain at which this proxy will be reached. e.g., netbird.example.com")
rootCmd.Flags().StringVar(&certDir, "cert-dir", envStringOrDefault("NB_PROXY_CERTIFICATE_DIRECTORY", "./certs"), "Directory to store certificates")
rootCmd.Flags().BoolVar(&acmeCerts, "acme-certs", envBoolOrDefault("NB_PROXY_ACME_CERTIFICATES", false), "Generate ACME certificates automatically")
rootCmd.Flags().StringVar(&acmeAddr, "acme-addr", envStringOrDefault("NB_PROXY_ACME_ADDRESS", ":80"), "HTTP address for ACME HTTP-01 challenges (only used when acme-challenge-type is http-01)")
rootCmd.Flags().StringVar(&acmeDir, "acme-dir", envStringOrDefault("NB_PROXY_ACME_DIRECTORY", acme.LetsEncryptURL), "URL of ACME challenge directory")
rootCmd.Flags().StringVar(&acmeEABKID, "acme-eab-kid", envStringOrDefault("NB_PROXY_ACME_EAB_KID", ""), "ACME EAB KID for account registration")
rootCmd.Flags().StringVar(&acmeEABHMACKey, "acme-eab-hmac-key", envStringOrDefault("NB_PROXY_ACME_EAB_HMAC_KEY", ""), "ACME EAB HMAC key for account registration")
rootCmd.Flags().StringVar(&acmeChallengeType, "acme-challenge-type", envStringOrDefault("NB_PROXY_ACME_CHALLENGE_TYPE", "tls-alpn-01"), "ACME challenge type: tls-alpn-01 (default, port 443 only) or http-01 (requires port 80)")
rootCmd.Flags().BoolVar(&debugEndpoint, "debug-endpoint", envBoolOrDefault("NB_PROXY_DEBUG_ENDPOINT", false), "Enable debug HTTP endpoint")
rootCmd.Flags().StringVar(&debugEndpointAddr, "debug-endpoint-addr", envStringOrDefault("NB_PROXY_DEBUG_ENDPOINT_ADDRESS", "localhost:8444"), "Address for the debug HTTP endpoint")
rootCmd.Flags().StringVar(&healthAddr, "health-addr", envStringOrDefault("NB_PROXY_HEALTH_ADDRESS", "localhost:8080"), "Address for the health probe endpoint (liveness/readiness/startup)")
rootCmd.Flags().StringVar(&forwardedProto, "forwarded-proto", envStringOrDefault("NB_PROXY_FORWARDED_PROTO", "auto"), "X-Forwarded-Proto value for backends: auto, http, or https")
rootCmd.Flags().StringVar(&trustedProxies, "trusted-proxies", envStringOrDefault("NB_PROXY_TRUSTED_PROXIES", ""), "Comma-separated list of trusted upstream proxy CIDR ranges (e.g. '10.0.0.0/8,192.168.1.1')")
rootCmd.Flags().StringVar(&certFile, "cert-file", envStringOrDefault("NB_PROXY_CERTIFICATE_FILE", "tls.crt"), "TLS certificate filename within the certificate directory")
rootCmd.Flags().StringVar(&certKeyFile, "cert-key-file", envStringOrDefault("NB_PROXY_CERTIFICATE_KEY_FILE", "tls.key"), "TLS certificate key filename within the certificate directory")
rootCmd.Flags().StringVar(&certLockMethod, "cert-lock-method", envStringOrDefault("NB_PROXY_CERT_LOCK_METHOD", "auto"), "Certificate lock method for cross-replica coordination: auto, flock, or k8s-lease")
rootCmd.Flags().StringVar(&wildcardCertDir, "wildcard-cert-dir", envStringOrDefault("NB_PROXY_WILDCARD_CERT_DIR", ""), "Directory containing wildcard certificate pairs (<name>.crt/<name>.key). Wildcard patterns are extracted from SANs automatically")
rootCmd.Flags().Uint16Var(&wgPort, "wg-port", envUint16OrDefault("NB_PROXY_WG_PORT", 0), "WireGuard listen port (0 = random). Fixed port only works with single-account deployments")
rootCmd.Flags().BoolVar(&proxyProtocol, "proxy-protocol", envBoolOrDefault("NB_PROXY_PROXY_PROTOCOL", false), "Enable PROXY protocol on TCP listeners to preserve client IPs behind L4 proxies")
rootCmd.Flags().StringVar(&preSharedKey, "preshared-key", envStringOrDefault("NB_PROXY_PRESHARED_KEY", ""), "Define a pre-shared key for the tunnel between proxy and peers")
rootCmd.Flags().BoolVar(&supportsCustomPorts, "supports-custom-ports", envBoolOrDefault("NB_PROXY_SUPPORTS_CUSTOM_PORTS", true), "Whether the proxy can bind arbitrary ports for UDP/TCP passthrough")
rootCmd.Flags().BoolVar(&requireSubdomain, "require-subdomain", envBoolOrDefault("NB_PROXY_REQUIRE_SUBDOMAIN", false), "Require a subdomain label in front of the cluster domain")
rootCmd.Flags().BoolVar(&private, "private", envBoolOrDefault("NB_PROXY_PRIVATE", false), "Enable private services accessible with NetBird-Only authentication mode.")
_ = rootCmd.Flags().MarkHidden("private")
rootCmd.Flags().DurationVar(&maxDialTimeout, "max-dial-timeout", envDurationOrDefault("NB_PROXY_MAX_DIAL_TIMEOUT", 0), "Cap per-service backend dial timeout (0 = no cap)")
rootCmd.Flags().DurationVar(&maxSessionIdleTimeout, "max-session-idle-timeout", envDurationOrDefault("NB_PROXY_MAX_SESSION_IDLE_TIMEOUT", 0), "Cap per-service session idle timeout (0 = no cap)")
rootCmd.Flags().StringVar(&geoDataDir, "geo-data-dir", envStringOrDefault("NB_PROXY_GEO_DATA_DIR", "/var/lib/netbird/geolocation"), "Directory for the GeoLite2 MMDB file (auto-downloaded if missing)")
rootCmd.Flags().StringVar(&crowdsecAPIURL, "crowdsec-api-url", envStringOrDefault("NB_PROXY_CROWDSEC_API_URL", ""), "CrowdSec LAPI URL for IP reputation checks")
rootCmd.Flags().StringVar(&crowdsecAPIKey, "crowdsec-api-key", envStringOrDefault("NB_PROXY_CROWDSEC_API_KEY", ""), "CrowdSec bouncer API key")
}
// Execute runs the root command.
func Execute() {
if err := rootCmd.Execute(); err != nil {
os.Exit(1)
}
}
// SetVersionInfo sets version information for the CLI.
func SetVersionInfo(version, commit, buildDate, goVersion string) {
Version = version
Commit = commit
BuildDate = buildDate
GoVersion = goVersion
rootCmd.Version = version
rootCmd.SetVersionTemplate("Version: {{.Version}}, Commit: " + Commit + ", BuildDate: " + BuildDate + ", Go: " + GoVersion + "\n")
}
func runServer(cmd *cobra.Command, args []string) error {
proxyToken := os.Getenv(envProxyToken)
if proxyToken == "" {
return fmt.Errorf("proxy token is required: set %s environment variable", envProxyToken)
}
level := logLevel
if debugLogs {
level = "debug"
}
logger := log.New()
_ = util.InitLogger(logger, level, util.LogConsole)
logger.Infof("configured log level: %s", level)
var wgPool, wgBatch uint64
var perf embed.Performance
if raw := os.Getenv(envPreallocatedBuffers); raw != "" {
n, err := strconv.ParseUint(raw, 10, 32)
if err != nil {
return fmt.Errorf("invalid %s %q: %w", envPreallocatedBuffers, raw, err)
}
wgPool = n
v := uint32(n)
perf.PreallocatedBuffersPerPool = &v
logger.Infof("tunnel preallocated buffers per pool: %d", n)
}
if raw := os.Getenv(envMaxBatchSize); raw != "" {
n, err := strconv.ParseUint(raw, 10, 32)
if err != nil {
return fmt.Errorf("invalid %s %q: %w", envMaxBatchSize, raw, err)
}
wgBatch = n
v := uint32(n)
perf.MaxBatchSize = &v
logger.Infof("tunnel max batch size override: %d", n)
}
if wgPool > 0 {
// Each bind recv goroutine (IPv4 + IPv6 + ICE relay) plus
// RoutineReadFromTUN eagerly reserves `batch` message buffers for
// the lifetime of the Device. A pool cap below that floor blocks
// the receive pipeline at startup.
batch := wgBatch
if batch == 0 {
batch = 128
}
const recvGoroutines = 4
floor := batch * recvGoroutines
if wgPool < floor {
logger.Warnf("%s=%d is below the eager-allocation floor (~%d for batch=%d); startup may deadlock",
envPreallocatedBuffers, wgPool, floor, batch)
}
}
switch forwardedProto {
case "auto", "http", "https":
default:
return fmt.Errorf("invalid --forwarded-proto value %q: must be auto, http, or https", forwardedProto)
}
_, err := domain.ValidateDomains([]string{proxyDomain})
if err != nil {
return fmt.Errorf("invalid domain value %q: %w", proxyDomain, err)
}
parsedTrustedProxies, err := trustedproxy.Parse(trustedProxies)
if err != nil {
return fmt.Errorf("invalid --trusted-proxies: %w", err)
}
ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGTERM, syscall.SIGINT)
defer stop()
srv := proxy.New(ctx, proxy.Config{
ListenAddr: addr,
Logger: logger,
Version: Version,
ManagementAddress: mgmtAddr,
ProxyURL: proxyDomain,
ProxyToken: proxyToken,
CertificateDirectory: certDir,
CertificateFile: certFile,
CertificateKeyFile: certKeyFile,
GenerateACMECertificates: acmeCerts,
ACMEChallengeAddress: acmeAddr,
ACMEDirectory: acmeDir,
ACMEEABKID: acmeEABKID,
ACMEEABHMACKey: acmeEABHMACKey,
ACMEChallengeType: acmeChallengeType,
DebugEndpointEnabled: debugEndpoint,
DebugEndpointAddress: debugEndpointAddr,
HealthAddr: healthAddr,
ForwardedProto: forwardedProto,
TrustedProxies: parsedTrustedProxies,
CertLockMethod: nbacme.CertLockMethod(certLockMethod),
WildcardCertDir: wildcardCertDir,
WireguardPort: wgPort,
Performance: perf,
ProxyProtocol: proxyProtocol,
PreSharedKey: preSharedKey,
SupportsCustomPorts: supportsCustomPorts,
RequireSubdomain: requireSubdomain,
Private: private,
MaxDialTimeout: maxDialTimeout,
MaxSessionIdleTimeout: maxSessionIdleTimeout,
MappingBatchWatchdog: envDurationOrDefault("NB_PROXY_MAPPING_BATCH_WATCHDOG", 0),
GeoDataDir: geoDataDir,
CrowdSecAPIURL: crowdsecAPIURL,
CrowdSecAPIKey: crowdsecAPIKey,
})
return srv.ListenAndServe(ctx, addr)
}
func envBoolOrDefault(key string, def bool) bool {
v, exists := os.LookupEnv(key)
if !exists {
return def
}
parsed, err := strconv.ParseBool(v)
if err != nil {
log.Warnf("parse %s=%q: %v, using default %v", key, v, err, def)
return def
}
return parsed
}
func envStringOrDefault(key string, def string) string {
v, exists := os.LookupEnv(key)
if !exists {
return def
}
return v
}
func envUint16OrDefault(key string, def uint16) uint16 {
v, exists := os.LookupEnv(key)
if !exists {
return def
}
parsed, err := strconv.ParseUint(v, 10, 16)
if err != nil {
log.Warnf("parse %s=%q: %v, using default %d", key, v, err, def)
return def
}
return uint16(parsed)
}
func envDurationOrDefault(key string, def time.Duration) time.Duration {
v, exists := os.LookupEnv(key)
if !exists {
return def
}
parsed, err := time.ParseDuration(v)
if err != nil {
log.Warnf("parse %s=%q: %v, using default %s", key, v, err, def)
return def
}
return parsed
}
proxy/cmd/proxy/main.go
+34
View File
@@ -0,0 +1,34 @@
package main
import (
"net/http"
// nolint:gosec
_ "net/http/pprof"
"runtime"
log "github.com/sirupsen/logrus"
"github.com/netbirdio/netbird/proxy/cmd/proxy/cmd"
)
var (
// Version is the application version (set via ldflags during build)
Version = "dev"
// Commit is the git commit hash (set via ldflags during build)
Commit = "unknown"
// BuildDate is the build date (set via ldflags during build)
BuildDate = "unknown"
// GoVersion is the Go version used to build the binary
GoVersion = runtime.Version()
)
func main() {
go func() {
log.Println(http.ListenAndServe("localhost:6060", nil))
}()
cmd.SetVersionInfo(Version, Commit, BuildDate, GoVersion)
cmd.Execute()
}
proxy/handle_mapping_stream_test.go
+104
View File
@@ -0,0 +1,104 @@
package proxy
import (
"context"
"io"
"testing"
"time"
log "github.com/sirupsen/logrus"
"github.com/stretchr/testify/assert"
"google.golang.org/grpc"
"google.golang.org/grpc/metadata"
"github.com/netbirdio/netbird/proxy/internal/health"
"github.com/netbirdio/netbird/shared/management/proto"
)
type mockMappingStream struct {
grpc.ClientStream
messages []*proto.GetMappingUpdateResponse
idx int
}
func (m *mockMappingStream) Recv() (*proto.GetMappingUpdateResponse, error) {
if m.idx >= len(m.messages) {
return nil, io.EOF
}
msg := m.messages[m.idx]
m.idx++
return msg, nil
}
func (m *mockMappingStream) Header() (metadata.MD, error) {
return nil, nil //nolint:nilnil
}
func (m *mockMappingStream) Trailer() metadata.MD { return nil }
func (m *mockMappingStream) CloseSend() error { return nil }
func (m *mockMappingStream) Context() context.Context { return context.Background() }
func (m *mockMappingStream) SendMsg(any) error { return nil }
func (m *mockMappingStream) RecvMsg(any) error { return nil }
func closedChan() chan struct{} {
ch := make(chan struct{})
close(ch)
return ch
}
func TestHandleMappingStream_SyncCompleteFlag(t *testing.T) {
checker := health.NewChecker(nil, nil)
s := &Server{
Logger: log.StandardLogger(),
healthChecker: checker,
routerReady: closedChan(),
}
stream := &mockMappingStream{
messages: []*proto.GetMappingUpdateResponse{
{InitialSyncComplete: true},
},
}
syncDone := false
err := s.handleMappingStream(context.Background(), stream, &syncDone, time.Time{})
assert.NoError(t, err)
assert.True(t, syncDone, "initial sync should be marked done when flag is set")
}
func TestHandleMappingStream_NoSyncFlagDoesNotMarkDone(t *testing.T) {
checker := health.NewChecker(nil, nil)
s := &Server{
Logger: log.StandardLogger(),
healthChecker: checker,
routerReady: closedChan(),
}
stream := &mockMappingStream{
messages: []*proto.GetMappingUpdateResponse{
{}, // no sync flag
},
}
syncDone := false
err := s.handleMappingStream(context.Background(), stream, &syncDone, time.Time{})
assert.NoError(t, err)
assert.False(t, syncDone, "initial sync should not be marked done without flag")
}
func TestHandleMappingStream_NilHealthChecker(t *testing.T) {
s := &Server{
Logger: log.StandardLogger(),
routerReady: closedChan(),
}
stream := &mockMappingStream{
messages: []*proto.GetMappingUpdateResponse{
{InitialSyncComplete: true},
},
}
syncDone := false
err := s.handleMappingStream(context.Background(), stream, &syncDone, time.Time{})
assert.NoError(t, err)
assert.True(t, syncDone, "sync done flag should be set even without health checker")
}
proxy/inbound.go
+572
View File
@@ -0,0 +1,572 @@
package proxy
import (
"context"
"crypto/tls"
"errors"
"fmt"
"io"
stdlog "log"
"net"
"net/http"
"net/netip"
"strconv"
"sync"
"time"
log "github.com/sirupsen/logrus"
"github.com/netbirdio/netbird/client/embed"
"github.com/netbirdio/netbird/proxy/internal/auth"
"github.com/netbirdio/netbird/proxy/internal/debug"
nbtcp "github.com/netbirdio/netbird/proxy/internal/tcp"
"github.com/netbirdio/netbird/proxy/internal/types"
)
// httpInboundReadHeaderTimeout matches the host-listener read header timeout
// so per-account http.Servers don't leak idle connections.
const httpInboundReadHeaderTimeout = 30 * time.Second
// httpInboundIdleTimeout caps idle keep-alive on per-account inbound HTTP
// servers; matches the host listener.
const httpInboundIdleTimeout = 90 * time.Second
// inboundShutdownTimeout caps how long a per-account http.Server gets to
// drain in-flight requests during teardown.
const inboundShutdownTimeout = 5 * time.Second
// privateInboundPortHTTPS is the WG-side TLS port. Each account's
// embedded netstack binds independently, so a fixed port is fine.
const privateInboundPortHTTPS = 443
// privateInboundPortHTTP is the WG-side plain-HTTP port.
const privateInboundPortHTTP = 80
// inboundManager wires per-account inbound listeners into the proxy
// pipeline when --private is enabled. When disabled the manager
// is nil and every method on *Server that touches it short-circuits.
type inboundManager struct {
logger *log.Logger
handler http.Handler
tlsConfig *tls.Config
// muxLock guards entries and pendingRoutes.
muxLock sync.Mutex
entries map[types.AccountID]*inboundEntry
pendingRoutes map[types.AccountID][]pendingInboundRoute
}
// inboundEntry owns the listeners, router and HTTP servers for a single
// account's embedded netstack. errorLogWriters retain the logrus pipe
// writers backing each http.Server's ErrorLog so tearDown can close
// them — otherwise the pipe + its scanner goroutine leak per account.
type inboundEntry struct {
router *nbtcp.Router
tlsListener net.Listener
plainListener net.Listener
httpsServer *http.Server
httpServer *http.Server
errorLogWriters []*io.PipeWriter
cancel context.CancelFunc
wg sync.WaitGroup
}
// pendingInboundRoute holds a route that arrived before the account's
// listener finished starting.
type pendingInboundRoute struct {
host nbtcp.SNIHost
route nbtcp.Route
}
// newInboundManager constructs a manager bound to the proxy's HTTP
// handler chain and TLS config.
func newInboundManager(logger *log.Logger, handler http.Handler, tlsConfig *tls.Config) *inboundManager {
return &inboundManager{
logger: logger,
handler: handler,
tlsConfig: tlsConfig,
entries: make(map[types.AccountID]*inboundEntry),
pendingRoutes: make(map[types.AccountID][]pendingInboundRoute),
}
}
// onClientReady is registered with NetBird.SetClientLifecycle so the
// listener pair comes up exactly when the embedded client reports ready.
// The returned value is opaque to the roundtrip package; it is handed
// back verbatim to onClientStop on teardown.
func (m *inboundManager) onClientReady(ctx context.Context, accountID types.AccountID, client *embed.Client) any {
if m == nil {
return nil
}
entry, err := m.bringUp(ctx, accountID, client)
if err != nil {
m.logger.WithField("account_id", accountID).WithError(err).Warn("failed to start per-account inbound listener; continuing without inbound")
return nil
}
m.flushPending(accountID, entry)
m.logger.WithFields(log.Fields{
"account_id": accountID,
"https": entry.tlsListener.Addr().String(),
"http": entry.plainListener.Addr().String(),
}).Info("per-account inbound listeners up")
return entry
}
// onClientStop tears down a per-account listener bundle. State is the
// opaque value previously returned by onClientReady.
func (m *inboundManager) onClientStop(accountID types.AccountID, state any) {
if m == nil {
return
}
entry, ok := state.(*inboundEntry)
if !ok || entry == nil {
return
}
m.tearDown(accountID, entry)
}
// bringUp opens both listeners on the account's netstack, builds the
// router, and starts the parallel HTTP servers.
func (m *inboundManager) bringUp(ctx context.Context, accountID types.AccountID, client *embed.Client) (*inboundEntry, error) {
tlsListener, err := client.ListenTCP(fmt.Sprintf(":%d", privateInboundPortHTTPS))
if err != nil {
return nil, fmt.Errorf("listen tls on netstack: %w", err)
}
plainListener, err := client.ListenTCP(fmt.Sprintf(":%d", privateInboundPortHTTP))
if err != nil {
_ = tlsListener.Close()
return nil, fmt.Errorf("listen plain on netstack: %w", err)
}
router := nbtcp.NewRouter(m.logger, accountDialResolver(accountID, client), tlsListener.Addr(), nbtcp.WithPlainHTTP(plainListener.Addr()))
scopedHandler := withTunnelLookup(m.handler, accountTunnelLookup(client))
// markOverlayOrigin stamps every connection accepted by an inbound
// listener with a context value middlewares can read to skip
// geo/CrowdSec checks (the source address is always inside the
// NetBird CGNAT range and won't match either dataset).
markOverlayOrigin := func(ctx context.Context, _ net.Conn) context.Context {
return types.WithOverlayOrigin(ctx)
}
httpsErrLog, httpsErrW := newInboundErrorLog(m.logger, "https", accountID)
httpErrLog, httpErrW := newInboundErrorLog(m.logger, "http", accountID)
httpsServer := &http.Server{
Handler: scopedHandler,
TLSConfig: m.tlsConfig,
ReadHeaderTimeout: httpInboundReadHeaderTimeout,
IdleTimeout: httpInboundIdleTimeout,
ErrorLog: httpsErrLog,
ConnContext: markOverlayOrigin,
}
httpServer := &http.Server{
Handler: scopedHandler,
ReadHeaderTimeout: httpInboundReadHeaderTimeout,
IdleTimeout: httpInboundIdleTimeout,
ErrorLog: httpErrLog,
ConnContext: markOverlayOrigin,
}
runCtx, cancel := context.WithCancel(ctx)
entry := &inboundEntry{
router: router,
tlsListener: tlsListener,
plainListener: plainListener,
httpsServer: httpsServer,
httpServer: httpServer,
errorLogWriters: []*io.PipeWriter{httpsErrW, httpErrW},
cancel: cancel,
}
entry.wg.Add(1)
go func() {
defer entry.wg.Done()
if err := router.Serve(runCtx, tlsListener); err != nil {
m.logger.WithField("account_id", accountID).Debugf("per-account router stopped: %v", err)
}
}()
entry.wg.Add(1)
go func() {
defer entry.wg.Done()
if err := httpsServer.ServeTLS(router.HTTPListener(), "", ""); err != nil && !errors.Is(err, http.ErrServerClosed) {
m.logger.WithField("account_id", accountID).Debugf("per-account https server stopped: %v", err)
}
}()
entry.wg.Add(1)
go func() {
defer entry.wg.Done()
if err := httpServer.Serve(router.HTTPListenerPlain()); err != nil && !errors.Is(err, http.ErrServerClosed) {
m.logger.WithField("account_id", accountID).Debugf("per-account http server stopped: %v", err)
}
}()
entry.wg.Add(1)
go func() {
defer entry.wg.Done()
feedRouterFromListener(runCtx, plainListener, router, m.logger, accountID)
}()
m.muxLock.Lock()
m.entries[accountID] = entry
m.muxLock.Unlock()
return entry, nil
}
// tearDown shuts every goroutine down and closes the netstack listeners.
func (m *inboundManager) tearDown(accountID types.AccountID, entry *inboundEntry) {
m.muxLock.Lock()
if m.entries[accountID] == entry {
delete(m.entries, accountID)
delete(m.pendingRoutes, accountID)
}
m.muxLock.Unlock()
entry.cancel()
shutdownCtx, cancel := context.WithTimeout(context.Background(), inboundShutdownTimeout)
defer cancel()
if err := entry.httpsServer.Shutdown(shutdownCtx); err != nil {
m.logger.Debugf("per-account https shutdown: %v", err)
}
if err := entry.httpServer.Shutdown(shutdownCtx); err != nil {
m.logger.Debugf("per-account http shutdown: %v", err)
}
if err := entry.tlsListener.Close(); err != nil {
m.logger.Debugf("close per-account tls listener: %v", err)
}
if err := entry.plainListener.Close(); err != nil {
m.logger.Debugf("close per-account plain listener: %v", err)
}
entry.wg.Wait()
// Close the ErrorLog pipes only after the http.Servers have fully
// stopped so any straggling stdlib write doesn't race with the
// close. Each writer also tears down the logrus scanner goroutine.
for _, w := range entry.errorLogWriters {
if err := w.Close(); err != nil {
m.logger.Debugf("close per-account inbound error log writer: %v", err)
}
}
}
// AddRoute records an SNI/host route on the account's per-account router.
// Routes registered before the listener is up are queued and replayed
// once startup completes.
func (m *inboundManager) AddRoute(accountID types.AccountID, host nbtcp.SNIHost, route nbtcp.Route) {
if m == nil {
return
}
m.muxLock.Lock()
entry, ok := m.entries[accountID]
if !ok {
m.queuePendingLocked(accountID, host, route)
m.muxLock.Unlock()
return
}
router := entry.router
m.muxLock.Unlock()
router.AddRoute(host, route)
}
// RemoveRoute drops a previously registered route. Safe to call when the
// listener is not yet up; queued copies are pruned in that case.
func (m *inboundManager) RemoveRoute(accountID types.AccountID, host nbtcp.SNIHost, svcID types.ServiceID) {
if m == nil {
return
}
m.muxLock.Lock()
m.dropPendingLocked(accountID, host, svcID)
entry, ok := m.entries[accountID]
if !ok {
m.muxLock.Unlock()
return
}
router := entry.router
m.muxLock.Unlock()
router.RemoveRoute(host, svcID)
}
// queuePendingLocked stores or upserts a pending route. Caller holds muxLock.
func (m *inboundManager) queuePendingLocked(accountID types.AccountID, host nbtcp.SNIHost, route nbtcp.Route) {
queued := m.pendingRoutes[accountID]
for i, pr := range queued {
if pr.host == host && pr.route.ServiceID == route.ServiceID {
queued[i] = pendingInboundRoute{host: host, route: route}
m.pendingRoutes[accountID] = queued
return
}
}
m.pendingRoutes[accountID] = append(queued, pendingInboundRoute{host: host, route: route})
}
// dropPendingLocked removes any queued route matching host/svcID.
// Caller holds muxLock.
func (m *inboundManager) dropPendingLocked(accountID types.AccountID, host nbtcp.SNIHost, svcID types.ServiceID) {
queued, ok := m.pendingRoutes[accountID]
if !ok {
return
}
filtered := queued[:0]
for _, pr := range queued {
if pr.host == host && pr.route.ServiceID == svcID {
continue
}
filtered = append(filtered, pr)
}
if len(filtered) == 0 {
delete(m.pendingRoutes, accountID)
return
}
m.pendingRoutes[accountID] = filtered
}
// flushPending applies all queued routes to a freshly-up router.
func (m *inboundManager) flushPending(accountID types.AccountID, entry *inboundEntry) {
m.muxLock.Lock()
queued := m.pendingRoutes[accountID]
delete(m.pendingRoutes, accountID)
m.muxLock.Unlock()
for _, pr := range queued {
entry.router.AddRoute(pr.host, pr.route)
}
}
// HasInbound reports whether the manager has a live listener for the account.
// Used by tests.
func (m *inboundManager) HasInbound(accountID types.AccountID) bool {
if m == nil {
return false
}
m.muxLock.Lock()
defer m.muxLock.Unlock()
_, ok := m.entries[accountID]
return ok
}
// PendingRouteCount reports the number of queued routes for the account.
// Used by tests.
func (m *inboundManager) PendingRouteCount(accountID types.AccountID) int {
if m == nil {
return 0
}
m.muxLock.Lock()
defer m.muxLock.Unlock()
return len(m.pendingRoutes[accountID])
}
// InboundListenerInfo describes the bound addresses of a single
// per-account inbound listener. Both addresses live on the embedded
// netstack of the account's WireGuard client and share the same tunnel IP.
type InboundListenerInfo struct {
TunnelIP string
HTTPSPort uint16
HTTPPort uint16
}
// ListenerInfo returns the inbound listener addresses for the given
// account, or ok=false when the account has no live listener. Used by
// the status-update RPC and the debug HTTP handler to surface inbound
// reachability to operators.
func (m *inboundManager) ListenerInfo(accountID types.AccountID) (InboundListenerInfo, bool) {
if m == nil {
return InboundListenerInfo{}, false
}
m.muxLock.Lock()
defer m.muxLock.Unlock()
entry, ok := m.entries[accountID]
if !ok || entry == nil {
return InboundListenerInfo{}, false
}
return listenerInfoFromEntry(entry), true
}
// Snapshot returns the inbound listener state for every account that has
// a live listener at call time. Empty when --private is off or
// no accounts have come up yet.
func (m *inboundManager) Snapshot() map[types.AccountID]InboundListenerInfo {
if m == nil {
return nil
}
m.muxLock.Lock()
defer m.muxLock.Unlock()
if len(m.entries) == 0 {
return nil
}
out := make(map[types.AccountID]InboundListenerInfo, len(m.entries))
for id, entry := range m.entries {
if entry == nil {
continue
}
out[id] = listenerInfoFromEntry(entry)
}
return out
}
// listenerInfoFromEntry extracts the tunnel IP and ports from a live
// per-account entry. Both listeners are bound on the same netstack so
// their host components match; we still pull the TLS host as the
// authoritative source.
func listenerInfoFromEntry(entry *inboundEntry) InboundListenerInfo {
info := InboundListenerInfo{HTTPSPort: privateInboundPortHTTPS, HTTPPort: privateInboundPortHTTP}
if entry.tlsListener != nil {
host, port := splitHostPort(entry.tlsListener.Addr())
info.TunnelIP = host
if port != 0 {
info.HTTPSPort = port
}
}
if entry.plainListener != nil {
host, port := splitHostPort(entry.plainListener.Addr())
if info.TunnelIP == "" {
info.TunnelIP = host
}
if port != 0 {
info.HTTPPort = port
}
}
return info
}
// splitHostPort extracts host and port from a net.Addr, returning the
// zero values when the address is missing or malformed.
func splitHostPort(addr net.Addr) (string, uint16) {
if addr == nil {
return "", 0
}
host, portStr, err := net.SplitHostPort(addr.String())
if err != nil {
return "", 0
}
if portStr == "" {
return host, 0
}
port, err := strconv.ParseUint(portStr, 10, 16)
if err != nil {
return host, 0
}
return host, uint16(port)
}
// feedRouterFromListener accepts on the plain-HTTP netstack listener and
// hands every connection to the account's router. The router peeks the
// first byte and dispatches to the plain-HTTP channel for non-TLS
// streams or the TLS channel for ClientHellos that arrive on :80.
func feedRouterFromListener(ctx context.Context, ln net.Listener, router *nbtcp.Router, logger *log.Logger, accountID types.AccountID) {
go func() {
<-ctx.Done()
_ = ln.Close()
}()
var backoff nbtcp.AcceptBackoff
for {
conn, err := ln.Accept()
if err != nil {
if ctx.Err() != nil || nbtcp.IsClosedListenerErr(err) {
return
}
logger.WithField("account_id", accountID).Debugf("plain inbound accept: %v; backing off", err)
if !backoff.Backoff(ctx) {
return
}
continue
}
backoff.Reset()
router.HandleConn(ctx, conn)
}
}
// accountDialResolver returns a DialResolver bound to a single account's
// embedded client. The router only ever serves traffic for that account
// so the supplied accountID is ignored at dial time.
func accountDialResolver(_ types.AccountID, client *embed.Client) nbtcp.DialResolver {
return func(_ types.AccountID) (types.DialContextFunc, error) {
return client.DialContext, nil
}
}
// accountTunnelLookup returns a TunnelLookupFunc backed by the embedded
// client's peerstore for a single account. Phase 3 uses the result to
// short-circuit ValidateTunnelPeer when the source IP is not in the
// account's roster and to seed the cached identity for known peers.
func accountTunnelLookup(client *embed.Client) auth.TunnelLookupFunc {
if client == nil {
return nil
}
return func(ip netip.Addr) (auth.PeerIdentity, bool) {
pubKey, fqdn, ok := client.IdentityForIP(ip)
if !ok {
return auth.PeerIdentity{}, false
}
return auth.PeerIdentity{
PubKey: pubKey,
TunnelIP: ip,
FQDN: fqdn,
}, true
}
}
// withTunnelLookup returns an http.Handler that attaches the per-account
// peerstore lookup to every request's context before delegating to next.
// Calling on the host-level listener is a no-op because that path never
// installs this wrapper, so the existing behaviour stays byte-for-byte
// identical when --private is off or the request didn't arrive
// on a per-account listener.
func withTunnelLookup(next http.Handler, lookup auth.TunnelLookupFunc) http.Handler {
if lookup == nil {
return next
}
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
ctx := auth.WithTunnelLookup(r.Context(), lookup)
next.ServeHTTP(w, r.WithContext(ctx))
})
}
// inboundDebugAdapter adapts *inboundManager to the debug.InboundProvider
// interface so the debug HTTP handler can render per-account inbound
// listener state without importing the proxy package.
type inboundDebugAdapter struct {
mgr *inboundManager
}
// InboundListeners returns a snapshot of the live per-account inbound
// listeners formatted for the debug surface.
func (a inboundDebugAdapter) InboundListeners() map[types.AccountID]debug.InboundListenerInfo {
if a.mgr == nil {
return nil
}
snap := a.mgr.Snapshot()
if len(snap) == 0 {
return nil
}
out := make(map[types.AccountID]debug.InboundListenerInfo, len(snap))
for id, info := range snap {
out[id] = debug.InboundListenerInfo{
TunnelIP: info.TunnelIP,
HTTPSPort: info.HTTPSPort,
HTTPPort: info.HTTPPort,
}
}
return out
}
// newInboundErrorLog routes a per-account http.Server's stdlib error
// stream through logrus at warn level. The returned PipeWriter must be
// closed by the caller (tearDown) once the http.Server has shut down —
// otherwise the pipe and its scanner goroutine leak per account, see
// logrus.Entry.WriterLevel.
func newInboundErrorLog(logger *log.Logger, scheme string, accountID types.AccountID) (*stdlog.Logger, *io.PipeWriter) {
w := logger.WithFields(log.Fields{
"inbound-http": scheme,
"account_id": accountID,
}).WriterLevel(log.WarnLevel)
return stdlog.New(w, "", 0), w
}
proxy/inbound_test.go
+657
View File
@@ -0,0 +1,657 @@
package proxy
import (
"bufio"
"context"
"crypto/tls"
"io"
"net"
"net/http"
"net/http/httptest"
"net/netip"
"sync"
"sync/atomic"
"testing"
"time"
log "github.com/sirupsen/logrus"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"google.golang.org/grpc"
"github.com/netbirdio/netbird/proxy/internal/auth"
"github.com/netbirdio/netbird/proxy/internal/roundtrip"
nbtcp "github.com/netbirdio/netbird/proxy/internal/tcp"
"github.com/netbirdio/netbird/proxy/internal/types"
"github.com/netbirdio/netbird/shared/management/proto"
)
// bufioReader wraps the connection in a buffered reader so http.ReadResponse
// can parse the response line + headers off the wire.
func bufioReader(conn net.Conn) *bufio.Reader {
return bufio.NewReader(conn)
}
// quietLogger returns a logger that emits nothing — keeps test output tidy.
func quietLogger() *log.Logger {
logger := log.New()
logger.SetLevel(log.PanicLevel)
return logger
}
func TestInboundManager_RouteScopedToAccount(t *testing.T) {
mgr := newInboundManager(quietLogger(), http.NotFoundHandler(), nil)
accountA := types.AccountID("acct-a")
accountB := types.AccountID("acct-b")
mgr.AddRoute(accountA, "shared.example", nbtcp.Route{Type: nbtcp.RouteHTTP, AccountID: accountA, ServiceID: "svc-a", Domain: "shared.example"})
mgr.AddRoute(accountB, "other.example", nbtcp.Route{Type: nbtcp.RouteHTTP, AccountID: accountB, ServiceID: "svc-b", Domain: "other.example"})
require.Equal(t, 1, mgr.PendingRouteCount(accountA), "account A should have one queued route")
require.Equal(t, 1, mgr.PendingRouteCount(accountB), "account B should have one queued route")
mgr.RemoveRoute(accountA, "shared.example", "svc-a")
mgr.RemoveRoute(accountB, "other.example", "svc-b")
assert.Equal(t, 0, mgr.PendingRouteCount(accountA), "queue should drain on remove")
assert.Equal(t, 0, mgr.PendingRouteCount(accountB), "queue should drain on remove")
}
func TestInboundManager_PendingThenFlush(t *testing.T) {
mgr := newInboundManager(quietLogger(), http.NotFoundHandler(), nil)
accountID := types.AccountID("acct-1")
host := nbtcp.SNIHost("example.test")
route := nbtcp.Route{Type: nbtcp.RouteHTTP, AccountID: accountID, ServiceID: "svc-1", Domain: "example.test"}
mgr.AddRoute(accountID, host, route)
require.Equal(t, 1, mgr.PendingRouteCount(accountID), "pending count before listener is up")
// Simulate listener up by registering a fake entry, then flushing.
router := nbtcp.NewRouter(quietLogger(), nil, &fakeAddr{addr: "127.0.0.1:0"})
entry := &inboundEntry{router: router}
mgr.muxLock.Lock()
mgr.entries[accountID] = entry
mgr.muxLock.Unlock()
mgr.flushPending(accountID, entry)
assert.Equal(t, 0, mgr.PendingRouteCount(accountID), "queue should be empty after flush")
}
// fakeAddr is a stub net.Addr for tests that don't actually bind sockets.
type fakeAddr struct {
addr string
}
func (a *fakeAddr) Network() string { return "tcp" }
func (a *fakeAddr) String() string { return a.addr }
// fakeMgmtClient implements roundtrip.managementClient for tests.
type fakeMgmtClient struct{}
func (fakeMgmtClient) CreateProxyPeer(_ context.Context, _ *proto.CreateProxyPeerRequest, _ ...grpc.CallOption) (*proto.CreateProxyPeerResponse, error) {
return &proto.CreateProxyPeerResponse{Success: true}, nil
}
// TestServer_PrivateInbound_NotEnabled_NoManager confirms that with
// --private off the inbound manager is nil and the standalone proxy
// keeps its zero-overhead default path.
func TestServer_PrivateInbound_NotEnabled_NoManager(t *testing.T) {
s := &Server{Logger: quietLogger(), Private: false}
s.initPrivateInbound(http.NotFoundHandler(), nil)
assert.Nil(t, s.inbound, "manager should remain nil when --private is off")
}
// TestServer_PrivateInbound_Enabled_WiresLifecycle confirms that
// --private alone wires the manager into the NetBird transport, so
// AddPeer / RemovePeer drive the lifecycle.
func TestServer_PrivateInbound_Enabled_WiresLifecycle(t *testing.T) {
s := &Server{Logger: quietLogger(), Private: true}
// Construct a NetBird transport. We can't actually start the embedded
// client here (that needs a real management server), but we can
// confirm that the lifecycle callbacks are registered.
s.netbird = roundtrip.NewNetBird(t.Context(), "test", "test", roundtrip.ClientConfig{
MgmtAddr: "http://invalid.test",
}, quietLogger(), nil, fakeMgmtClient{})
s.initPrivateInbound(http.NotFoundHandler(), &tls.Config{}) //nolint:gosec
require.NotNil(t, s.inbound, "manager should be set when --private is on")
assert.NotNil(t, s.inbound.handler, "handler should be set on manager")
assert.NotNil(t, s.inbound.tlsConfig, "tls config should be set on manager")
}
// TestInboundManager_AddRouteAfterReady_RegistersDirectly verifies that
// when the listener is already up, AddRoute writes straight to the
// router without queueing.
func TestInboundManager_AddRouteAfterReady_RegistersDirectly(t *testing.T) {
mgr := newInboundManager(quietLogger(), http.NotFoundHandler(), nil)
accountID := types.AccountID("acct-1")
router := nbtcp.NewRouter(quietLogger(), nil, &fakeAddr{addr: "127.0.0.1:0"})
mgr.muxLock.Lock()
mgr.entries[accountID] = &inboundEntry{router: router}
mgr.muxLock.Unlock()
host := nbtcp.SNIHost("ready.example")
mgr.AddRoute(accountID, host, nbtcp.Route{Type: nbtcp.RouteHTTP, AccountID: accountID, ServiceID: "svc-ready", Domain: string(host)})
assert.Equal(t, 0, mgr.PendingRouteCount(accountID), "no pending entries when listener is up")
}
// TestPrivateCapability_DerivedFromPrivateOnly tests that the capability
// bit reported upstream tracks --private exclusively. The previous
// --private flag has been folded into --private.
func TestPrivateCapability_DerivedFromPrivateOnly(t *testing.T) {
tests := []struct {
name string
private bool
expected bool
}{
{"off", false, false},
{"on", true, true},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
s := &Server{Private: tt.private}
assert.Equal(t, tt.expected, s.Private, "private capability bit should match --private")
})
}
}
// TestInboundManager_RouteScopedToAccountB_DoesNotMatchA verifies that a
// service registered for account B is invisible to a router serving
// account A. We exercise the path through real per-account routers.
func TestInboundManager_RouteScopedToAccountB_DoesNotMatchA(t *testing.T) {
mgr := newInboundManager(quietLogger(), http.NotFoundHandler(), nil)
accountA := types.AccountID("acct-a")
accountB := types.AccountID("acct-b")
routerA := nbtcp.NewRouter(quietLogger(), nil, &fakeAddr{addr: "127.0.0.1:0"})
routerB := nbtcp.NewRouter(quietLogger(), nil, &fakeAddr{addr: "127.0.0.1:0"})
mgr.muxLock.Lock()
mgr.entries[accountA] = &inboundEntry{router: routerA}
mgr.entries[accountB] = &inboundEntry{router: routerB}
mgr.muxLock.Unlock()
host := nbtcp.SNIHost("shared.example")
mgr.AddRoute(accountB, host, nbtcp.Route{Type: nbtcp.RouteHTTP, AccountID: accountB, ServiceID: "svc-b", Domain: string(host)})
// Account A's router should have no routes; account B's should have one.
// We check via IsEmpty — true means no routes and no fallback.
assert.True(t, routerA.IsEmpty(), "account A router must not see account B's mappings")
assert.False(t, routerB.IsEmpty(), "account B router should hold its own mapping")
}
// TestInboundEntry_ShutdownIdempotent ensures that tearDown can run twice
// without panicking — callers may invoke it from RemovePeer + StopAll.
func TestInboundEntry_ShutdownIdempotent(t *testing.T) {
t.Skip("teardown requires real netstack listeners; covered by integration tests")
}
// TestRouter_PlainHTTP_ForwardedProtoIsHTTP exercises the full per-account
// router pipeline against a loopback listener (proxy of a netstack
// listener for test purposes): a plain HTTP request lands on the plain
// http.Server and the inner handler observes a nil r.TLS, which is what
// auth.ResolveProto translates to "http" in the real pipeline.
func TestRouter_PlainHTTP_ForwardedProtoIsHTTP(t *testing.T) {
logger := quietLogger()
var captured atomic.Value
captured.Store("")
handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.TLS == nil {
captured.Store("http")
} else {
captured.Store("https")
}
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte("ok"))
})
hostListener, err := net.Listen("tcp", "127.0.0.1:0")
require.NoError(t, err, "loopback listener bind must succeed")
defer hostListener.Close()
router := nbtcp.NewRouter(logger, nil, hostListener.Addr(), nbtcp.WithPlainHTTP(hostListener.Addr()))
httpServer := &http.Server{Handler: handler, ReadHeaderTimeout: time.Second}
defer func() { _ = httpServer.Close() }()
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
go func() { _ = httpServer.Serve(router.HTTPListenerPlain()) }()
go func() { _ = router.Serve(ctx, hostListener) }()
conn, err := net.DialTimeout("tcp", hostListener.Addr().String(), 2*time.Second)
require.NoError(t, err, "plain HTTP dial must succeed")
defer conn.Close()
_, err = conn.Write([]byte("GET / HTTP/1.1\r\nHost: example.com\r\nConnection: close\r\n\r\n"))
require.NoError(t, err, "write must succeed")
resp, err := http.ReadResponse(bufioReader(conn), nil)
require.NoError(t, err, "must read response")
defer resp.Body.Close()
assert.Equal(t, http.StatusOK, resp.StatusCode)
assert.Equal(t, "http", captured.Load(), "ForwardedProto must be http on plain path")
}
// TestWithTunnelLookup_AttachesLookupToContext verifies that requests
// flowing through the per-account handler wrapper carry the peerstore
// lookup function. Phase 3's local-first deny path depends on this.
func TestWithTunnelLookup_AttachesLookupToContext(t *testing.T) {
expected := auth.PeerIdentity{TunnelIP: netip.MustParseAddr("100.64.0.10"), FQDN: "peer.netbird"}
lookup := auth.TunnelLookupFunc(func(_ netip.Addr) (auth.PeerIdentity, bool) {
return expected, true
})
var observed auth.TunnelLookupFunc
inner := http.HandlerFunc(func(_ http.ResponseWriter, r *http.Request) {
observed = auth.TunnelLookupFromContext(r.Context())
})
handler := withTunnelLookup(inner, lookup)
r := httptest.NewRequest(http.MethodGet, "https://svc.example/", nil)
handler.ServeHTTP(httptest.NewRecorder(), r)
require.NotNil(t, observed, "wrapper must inject the lookup into the request context")
got, ok := observed(netip.MustParseAddr("100.64.0.10"))
assert.True(t, ok, "lookup must round-trip through context")
assert.Equal(t, expected.FQDN, got.FQDN, "lookup must return the same identity it was constructed with")
}
// TestWithTunnelLookup_NilLookupIsNoop confirms the wrapper is a pure
// pass-through when no lookup is provided. Required for the host-level
// listener path to keep its byte-for-byte previous behaviour.
func TestWithTunnelLookup_NilLookupIsNoop(t *testing.T) {
var called bool
inner := http.HandlerFunc(func(_ http.ResponseWriter, r *http.Request) {
called = true
assert.Nil(t, auth.TunnelLookupFromContext(r.Context()), "host-level path must not see a lookup function")
})
handler := withTunnelLookup(inner, nil)
r := httptest.NewRequest(http.MethodGet, "https://svc.example/", nil)
handler.ServeHTTP(httptest.NewRecorder(), r)
assert.True(t, called, "wrapper without lookup must still invoke next")
}
// fakeListener satisfies net.Listener for snapshot tests without binding
// a real socket on the netstack.
type fakeListener struct {
addr net.Addr
}
func (f *fakeListener) Accept() (net.Conn, error) { return nil, net.ErrClosed }
func (f *fakeListener) Close() error { return nil }
func (f *fakeListener) Addr() net.Addr { return f.addr }
// TestInboundManager_ListenerInfo confirms ListenerInfo and Snapshot
// surface the bound tunnel-IP and ports for live entries.
func TestInboundManager_ListenerInfo(t *testing.T) {
mgr := newInboundManager(quietLogger(), http.NotFoundHandler(), nil)
accountID := types.AccountID("acct-info")
tlsAddr := &net.TCPAddr{IP: net.ParseIP("100.64.0.5"), Port: privateInboundPortHTTPS}
plainAddr := &net.TCPAddr{IP: net.ParseIP("100.64.0.5"), Port: privateInboundPortHTTP}
mgr.muxLock.Lock()
mgr.entries[accountID] = &inboundEntry{
tlsListener: &fakeListener{addr: tlsAddr},
plainListener: &fakeListener{addr: plainAddr},
}
mgr.muxLock.Unlock()
info, ok := mgr.ListenerInfo(accountID)
require.True(t, ok, "ListenerInfo must report ok for live entry")
assert.Equal(t, "100.64.0.5", info.TunnelIP, "tunnel IP must come from listener address")
assert.Equal(t, uint16(privateInboundPortHTTPS), info.HTTPSPort, "TLS port must match bound port")
assert.Equal(t, uint16(privateInboundPortHTTP), info.HTTPPort, "HTTP port must match bound port")
snap := mgr.Snapshot()
require.Len(t, snap, 1, "snapshot must contain exactly one entry")
assert.Equal(t, info, snap[accountID], "snapshot entry must equal direct lookup")
_, ok = mgr.ListenerInfo(types.AccountID("missing"))
assert.False(t, ok, "ListenerInfo must report ok=false for unknown accounts")
}
// TestInboundManager_NilManagerSafe ensures the observability accessors
// are safe to call when --private is off (nil manager).
func TestInboundManager_NilManagerSafe(t *testing.T) {
var mgr *inboundManager
_, ok := mgr.ListenerInfo("anything")
assert.False(t, ok, "nil manager must return ok=false")
assert.Nil(t, mgr.Snapshot(), "nil manager must return nil snapshot")
}
// TestInboundManager_ConcurrentAddRemove pounds AddRoute / RemoveRoute
// from multiple goroutines to expose any locking gaps.
func TestInboundManager_ConcurrentAddRemove(t *testing.T) {
mgr := newInboundManager(quietLogger(), http.NotFoundHandler(), nil)
accountID := types.AccountID("acct-1")
const workers = 32
const iterations = 50
var wg sync.WaitGroup
wg.Add(workers)
for i := 0; i < workers; i++ {
go func(idx int) {
defer wg.Done()
host := nbtcp.SNIHost("example.test")
svc := types.ServiceID("svc")
route := nbtcp.Route{Type: nbtcp.RouteHTTP, AccountID: accountID, ServiceID: svc, Domain: "example.test"}
for j := 0; j < iterations; j++ {
mgr.AddRoute(accountID, host, route)
mgr.RemoveRoute(accountID, host, svc)
}
}(i)
}
done := make(chan struct{})
go func() {
wg.Wait()
close(done)
}()
select {
case <-done:
case <-time.After(10 * time.Second):
t.Fatal("concurrent add/remove timed out")
}
}
// TestFeedRouterFromListener_DeliversConnectionToHandler validates the
// per-account inbound chain end-to-end with a loopback listener
// substituted for the embedded netstack: a TCP connection arriving at
// the plain listener flows through feedRouterFromListener, the router's
// peek-and-dispatch, the wrapped HTTP server, and reaches the user
// handler. If the embedded netstack is delivering connections at all,
// this is the path they take. Failures localise to wiring bugs in the
// proxy, not the netstack.
func TestFeedRouterFromListener_DeliversConnectionToHandler(t *testing.T) {
logger := quietLogger()
hits := make(chan string, 1)
handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
hits <- r.Host
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte("served"))
})
plainLn, err := net.Listen("tcp", "127.0.0.1:0")
require.NoError(t, err, "plain loopback bind must succeed")
t.Cleanup(func() { _ = plainLn.Close() })
router := nbtcp.NewRouter(logger, nil, &fakeAddr{addr: "127.0.0.1:0"}, nbtcp.WithPlainHTTP(plainLn.Addr()))
httpServer := &http.Server{Handler: handler, ReadHeaderTimeout: time.Second}
t.Cleanup(func() { _ = httpServer.Close() })
ctx, cancel := context.WithCancel(context.Background())
t.Cleanup(cancel)
go func() { _ = httpServer.Serve(router.HTTPListenerPlain()) }()
go feedRouterFromListener(ctx, plainLn, router, logger, types.AccountID("acct-1"))
conn, err := net.DialTimeout("tcp", plainLn.Addr().String(), 2*time.Second)
require.NoError(t, err, "must connect to the plain listener")
t.Cleanup(func() { _ = conn.Close() })
_, err = conn.Write([]byte("GET / HTTP/1.1\r\nHost: app.example\r\nConnection: close\r\n\r\n"))
require.NoError(t, err, "request write must succeed")
resp, err := http.ReadResponse(bufioReader(conn), nil)
require.NoError(t, err, "must read response from server")
t.Cleanup(func() { _ = resp.Body.Close() })
assert.Equal(t, http.StatusOK, resp.StatusCode, "handler must be reached")
select {
case host := <-hits:
assert.Equal(t, "app.example", host, "handler must observe the request Host")
case <-time.After(2 * time.Second):
t.Fatal("handler was not invoked — connection did not flow through router → http server")
}
}
// TestFeedRouterFromListener_DispatchesTLSToTLSChannel verifies that a
// TLS ClientHello arriving on the plain listener is detected by the
// router peek and re-dispatched to the TLS channel — the cross-channel
// fallback the inbound stack relies on for HTTPS-on-:80 testing.
func TestFeedRouterFromListener_DispatchesTLSToTLSChannel(t *testing.T) {
logger := quietLogger()
hits := make(chan string, 1)
tlsHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
hits <- r.Host
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte("served-tls"))
})
plainLn, err := net.Listen("tcp", "127.0.0.1:0")
require.NoError(t, err, "plain loopback bind must succeed")
t.Cleanup(func() { _ = plainLn.Close() })
tlsLn, err := net.Listen("tcp", "127.0.0.1:0")
require.NoError(t, err, "tls loopback bind must succeed")
t.Cleanup(func() { _ = tlsLn.Close() })
router := nbtcp.NewRouter(logger, nil, tlsLn.Addr(), nbtcp.WithPlainHTTP(plainLn.Addr()))
tlsConfig := selfSignedTLSConfig(t)
httpsServer := &http.Server{
Handler: tlsHandler,
TLSConfig: tlsConfig,
ReadHeaderTimeout: time.Second,
}
t.Cleanup(func() { _ = httpsServer.Close() })
ctx, cancel := context.WithCancel(context.Background())
t.Cleanup(cancel)
go func() { _ = httpsServer.ServeTLS(router.HTTPListener(), "", "") }()
go feedRouterFromListener(ctx, plainLn, router, logger, types.AccountID("acct-tls"))
tlsConn, err := tls.Dial("tcp", plainLn.Addr().String(), &tls.Config{InsecureSkipVerify: true}) //nolint:gosec
require.NoError(t, err, "TLS dial against the plain listener must succeed (cross-channel)")
t.Cleanup(func() { _ = tlsConn.Close() })
req, err := http.NewRequest(http.MethodGet, "https://app.example/", nil)
require.NoError(t, err)
require.NoError(t, req.Write(tlsConn), "TLS request write must succeed")
resp, err := http.ReadResponse(bufioReader(tlsConn), req)
require.NoError(t, err, "must read TLS response")
t.Cleanup(func() { _ = resp.Body.Close() })
assert.Equal(t, http.StatusOK, resp.StatusCode, "TLS handler must be reached")
select {
case host := <-hits:
assert.Equal(t, "app.example", host, "TLS handler must observe the request Host")
case <-time.After(2 * time.Second):
t.Fatal("TLS handler was not invoked — peek/dispatch path is broken")
}
}
func selfSignedTLSConfig(t *testing.T) *tls.Config {
t.Helper()
cert, err := tls.X509KeyPair(testCertPEM, testKeyPEM)
require.NoError(t, err, "load static self-signed cert")
return &tls.Config{Certificates: []tls.Certificate{cert}, MinVersion: tls.VersionTLS12} //nolint:gosec
}
// TestNewInboundErrorLog_WriterIsCloseable guards the close path on the
// logrus PipeWriter that backs each per-account http.Server's ErrorLog.
// logrus.Entry.WriterLevel returns an *io.PipeWriter that owns a pipe +
// scanner goroutine; the caller must Close() it on teardown or the
// resources leak per account. The contract is verified two ways:
//
// - the constructor returns a non-nil writer the caller can keep,
// - writing to the writer after Close() fails with io.ErrClosedPipe,
// which is the only externally observable sign that Close was wired.
//
// A leaking refactor (forgetting to thread the writer to tearDown, or
// dropping the Close call) would still pass this test individually but
// fail an integration goleak check; this unit test is the cheap first
// line of defence.
func TestNewInboundErrorLog_WriterIsCloseable(t *testing.T) {
logger := quietLogger()
stdLog, writer := newInboundErrorLog(logger, "https", types.AccountID("acct-1"))
require.NotNil(t, stdLog, "newInboundErrorLog must return a non-nil *log.Logger")
require.NotNil(t, writer, "newInboundErrorLog must return the underlying PipeWriter so tearDown can Close it")
// First Close succeeds.
require.NoError(t, writer.Close(), "PipeWriter.Close should succeed the first time")
// After Close, the writer must refuse new writes — that's the only
// behavioural signal that the pipe (and its scanner goroutine) has
// shut down.
_, err := writer.Write([]byte("post-close write\n"))
require.ErrorIs(t, err, io.ErrClosedPipe,
"writes after Close must surface io.ErrClosedPipe so callers know the writer is gone")
}
// testCertPEM / testKeyPEM are a minimal RSA self-signed cert for
// 127.0.0.1 — only used by tests that need a working TLS handshake.
var testCertPEM = []byte(`-----BEGIN CERTIFICATE-----
MIIBhTCCASugAwIBAgIQIRi6zePL6mKjOipn+dNuaTAKBggqhkjOPQQDAjASMRAw
DgYDVQQKEwdBY21lIENvMB4XDTE3MTAyMDE5NDMwNloXDTE4MTAyMDE5NDMwNlow
EjEQMA4GA1UEChMHQWNtZSBDbzBZMBMGByqGSM49AgEGCCqGSM49AwEHA0IABD0d
7VNhbWvZLWPuj/RtHFjvtJBEwOkhbN/BnnE8rnZR8+sbwnc/KhCk3FhnpHZnQz7B
5aETbbIgmuvewdjvSBSjYzBhMA4GA1UdDwEB/wQEAwICpDATBgNVHSUEDDAKBggr
BgEFBQcDATAPBgNVHRMBAf8EBTADAQH/MCkGA1UdEQQiMCCCDmxvY2FsaG9zdDo1
NDUzgg4xMjcuMC4wLjE6NTQ1MzAKBggqhkjOPQQDAgNIADBFAiEA2zpJEPQyz6/l
Wf86aX6PepsntZv2GYlA5UpabfT2EZICICpJ5h/iI+i341gBmLiAFQOyTDT+/wQc
6MF9+Yw1Yy0t
-----END CERTIFICATE-----`)
var testKeyPEM = []byte(`-----BEGIN EC PRIVATE KEY-----
MHcCAQEEIIrYSSNQFaA2Hwf1duRSxKtLYX5CB04fSeQ6tF1aY/PuoAoGCCqGSM49
AwEHoUQDQgAEPR3tU2Fta9ktY+6P9G0cWO+0kETA6SFs38GecTyudlHz6xvCdz8q
EKTcWGekdmdDPsHloRNtsiCa697B2O9IFA==
-----END EC PRIVATE KEY-----`)
// scriptedAcceptListener returns pre-scripted errors from Accept(). Used
// to drive the feedRouterFromListener tests without binding a real
// socket — the production code path is a netstack-backed listener that
// returns gVisor's "endpoint is in invalid state" forever after its
// endpoint is destroyed.
type scriptedAcceptListener struct {
errs chan error
closed chan struct{}
}
func newScriptedAcceptListener(errs ...error) *scriptedAcceptListener {
s := &scriptedAcceptListener{
errs: make(chan error, len(errs)+1),
closed: make(chan struct{}),
}
for _, e := range errs {
s.errs <- e
}
return s
}
func (s *scriptedAcceptListener) Accept() (net.Conn, error) {
select {
case <-s.closed:
return nil, net.ErrClosed
case err := <-s.errs:
return nil, err
}
}
func (s *scriptedAcceptListener) Close() error {
select {
case <-s.closed:
default:
close(s.closed)
}
return nil
}
func (s *scriptedAcceptListener) Addr() net.Addr {
return &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 0}
}
// errSentinel carries a literal error message so tests can synthesise
// the exact gVisor text without importing the netstack package.
type errSentinel string
func (e errSentinel) Error() string { return string(e) }
// TestFeedRouterFromListener_ExitsOnGVisorInvalidEndpoint is the
// regression guard for the inbound side of the tight-loop bug. The
// per-account plain-HTTP feeder must recognise gVisor's "endpoint is in
// invalid state" and exit, otherwise it pegs a CPU core and floods the
// account-scoped log with the same accept error every iteration.
func TestFeedRouterFromListener_ExitsOnGVisorInvalidEndpoint(t *testing.T) {
logger := log.StandardLogger()
addr := &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 80}
router := nbtcp.NewRouter(logger, nil, addr)
gvisorErr := &net.OpError{
Op: "accept",
Net: "tcp",
Addr: addr,
Err: errSentinel("endpoint is in invalid state"),
}
ln := newScriptedAcceptListener(gvisorErr)
defer ln.Close()
done := make(chan struct{})
go func() {
defer close(done)
feedRouterFromListener(context.Background(), ln, router, logger, "acct-1")
}()
select {
case <-done:
// Expected: loop recognised the gVisor error and returned.
case <-time.After(2 * time.Second):
t.Fatal("feedRouterFromListener did not exit on gVisor 'endpoint is in invalid state' — accept loop is spinning")
}
}
// TestFeedRouterFromListener_BacksOffOnTransientError asserts the
// defence-in-depth path: an unknown sticky Accept error must NOT cause
// CPU spin. The loop backs off and exits cleanly when ctx is cancelled.
func TestFeedRouterFromListener_BacksOffOnTransientError(t *testing.T) {
logger := log.StandardLogger()
addr := &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 80}
router := nbtcp.NewRouter(logger, nil, addr)
const transientCount = 5
errs := make([]error, transientCount)
for i := range errs {
errs[i] = errSentinel("transient: temporary network error")
}
ln := newScriptedAcceptListener(errs...)
defer ln.Close()
ctx, cancel := context.WithCancel(context.Background())
start := time.Now()
done := make(chan struct{})
go func() {
defer close(done)
feedRouterFromListener(ctx, ln, router, logger, "acct-1")
}()
time.AfterFunc(150*time.Millisecond, cancel)
select {
case <-done:
// Expected.
case <-time.After(2 * time.Second):
t.Fatal("feedRouterFromListener did not exit on ctx cancellation — backoff or exit path broken")
}
// Without backoff the 5 scripted errors would burn in microseconds.
// With backoff the first delay alone is 5ms, so the loop must take
// at least that long even though ctx fires at 150ms.
elapsed := time.Since(start)
assert.GreaterOrEqual(t, elapsed, 5*time.Millisecond,
"loop ran without backing off — would burn CPU in production")
}
proxy/internal/accesslog/logger.go
+442
View File
@@ -0,0 +1,442 @@
package accesslog
import (
"context"
"maps"
"net/netip"
"sync"
"sync/atomic"
"time"
"github.com/rs/xid"
log "github.com/sirupsen/logrus"
"google.golang.org/grpc"
"google.golang.org/protobuf/types/known/timestamppb"
"github.com/netbirdio/netbird/proxy/auth"
"github.com/netbirdio/netbird/proxy/internal/types"
"github.com/netbirdio/netbird/shared/management/proto"
"github.com/netbirdio/netbird/trustedproxy"
)
const (
requestThreshold = 10000 // Log every 10k requests
bytesThreshold = 1024 * 1024 * 1024 // Log every 1GB
usageCleanupPeriod = 1 * time.Hour // Clean up stale counters every hour
usageInactiveWindow = 24 * time.Hour // Consider domain inactive if no traffic for 24 hours
logSendTimeout = 10 * time.Second
// denyCooldown is the min interval between deny log entries per service+reason
// to prevent flooding from denied connections (e.g. UDP packets from blocked IPs).
denyCooldown = 10 * time.Second
// maxDenyBuckets caps tracked deny rate-limit entries to bound memory under DDoS.
maxDenyBuckets = 10000
// maxLogWorkers caps concurrent gRPC send goroutines.
maxLogWorkers = 4096
)
type domainUsage struct {
requestCount int64
requestStartTime time.Time
bytesTransferred int64
bytesStartTime time.Time
lastActivity time.Time // Track last activity for cleanup
}
type gRPCClient interface {
SendAccessLog(ctx context.Context, in *proto.SendAccessLogRequest, opts ...grpc.CallOption) (*proto.SendAccessLogResponse, error)
}
// denyBucketKey identifies a rate-limited deny log stream.
type denyBucketKey struct {
ServiceID types.ServiceID
Reason string
}
// denyBucket tracks rate-limited deny log entries.
type denyBucket struct {
lastLogged time.Time
suppressed int64
}
// Logger sends access log entries to the management server via gRPC.
type Logger struct {
client gRPCClient
logger *log.Logger
trustedProxies *trustedproxy.List
usageMux sync.Mutex
domainUsage map[string]*domainUsage
denyMu sync.Mutex
denyBuckets map[denyBucketKey]*denyBucket
logSem chan struct{}
cleanupCancel context.CancelFunc
dropped atomic.Int64
}
// NewLogger creates a new access log Logger. The trustedProxies parameter
// configures which upstream proxy IP ranges are trusted for extracting
// the real client IP from X-Forwarded-For headers.
func NewLogger(client gRPCClient, logger *log.Logger, trustedProxies *trustedproxy.List) *Logger {
if logger == nil {
logger = log.StandardLogger()
}
ctx, cancel := context.WithCancel(context.Background())
l := &Logger{
client: client,
logger: logger,
trustedProxies: trustedProxies,
domainUsage: make(map[string]*domainUsage),
denyBuckets: make(map[denyBucketKey]*denyBucket),
logSem: make(chan struct{}, maxLogWorkers),
cleanupCancel: cancel,
}
// Start background cleanup routine
go l.cleanupStaleUsage(ctx)
return l
}
// Close stops the cleanup routine. Should be called during graceful shutdown.
func (l *Logger) Close() {
if l.cleanupCancel != nil {
l.cleanupCancel()
}
}
type logEntry struct {
ID string
AccountID types.AccountID
ServiceID types.ServiceID
Host string
Path string
DurationMs int64
Method string
ResponseCode int32
SourceIP netip.Addr
AuthMechanism string
UserID string
AuthSuccess bool
BytesUpload int64
BytesDownload int64
Protocol Protocol
Metadata map[string]string
AgentNetwork bool
}
// Protocol identifies the transport protocol of an access log entry.
type Protocol string
const (
ProtocolHTTP Protocol = "http"
ProtocolTCP Protocol = "tcp"
ProtocolUDP Protocol = "udp"
ProtocolTLS Protocol = "tls"
)
// L4Entry holds the data for a layer-4 (TCP/UDP) access log entry.
type L4Entry struct {
AccountID types.AccountID
ServiceID types.ServiceID
Protocol Protocol
Host string // SNI hostname or listen address
SourceIP netip.Addr
DurationMs int64
BytesUpload int64
BytesDownload int64
// DenyReason, when non-empty, indicates the connection was denied.
// Values match the HTTP auth mechanism strings: "ip_restricted",
// "country_restricted", "geo_unavailable", "crowdsec_ban", etc.
DenyReason string
// Metadata carries extra context about the connection (e.g. CrowdSec verdict).
Metadata map[string]string
}
// LogL4 sends an access log entry for a layer-4 connection (TCP or UDP).
// The call is non-blocking: the gRPC send happens in a background goroutine.
func (l *Logger) LogL4(entry L4Entry) {
le := logEntry{
ID: xid.New().String(),
AccountID: entry.AccountID,
ServiceID: entry.ServiceID,
Protocol: entry.Protocol,
Host: entry.Host,
SourceIP: entry.SourceIP,
DurationMs: entry.DurationMs,
BytesUpload: entry.BytesUpload,
BytesDownload: entry.BytesDownload,
Metadata: maps.Clone(entry.Metadata),
}
if entry.DenyReason != "" {
if !l.allowDenyLog(entry.ServiceID, entry.DenyReason) {
return
}
le.AuthMechanism = entry.DenyReason
le.AuthSuccess = false
}
l.log(le)
l.trackUsage(entry.Host, entry.BytesUpload+entry.BytesDownload)
}
// allowDenyLog rate-limits deny log entries per service+reason combination.
func (l *Logger) allowDenyLog(serviceID types.ServiceID, reason string) bool {
key := denyBucketKey{ServiceID: serviceID, Reason: reason}
now := time.Now()
l.denyMu.Lock()
defer l.denyMu.Unlock()
b, ok := l.denyBuckets[key]
if !ok {
if len(l.denyBuckets) >= maxDenyBuckets {
return false
}
l.denyBuckets[key] = &denyBucket{lastLogged: now}
return true
}
if now.Sub(b.lastLogged) >= denyCooldown {
if b.suppressed > 0 {
l.logger.Debugf("access restriction: suppressed %d deny log entries for %s (%s)", b.suppressed, serviceID, reason)
}
b.lastLogged = now
b.suppressed = 0
return true
}
b.suppressed++
return false
}
// usageMetadataKeys is the allowlist of metadata retained on a stripped,
// usage-only agent-network entry. Mirrors the llm.* / cost.* keys in
// proxy/internal/middleware/keys.go — only the dimensions management needs to
// record a usage row (provider / model / tokens / cost / groups).
var usageMetadataKeys = map[string]struct{}{
"llm.provider": {},
"llm.model": {},
"llm.resolved_provider_id": {},
"llm.input_tokens": {},
"llm.output_tokens": {},
"llm.total_tokens": {},
"llm.cached_input_tokens": {},
"llm.cache_creation_tokens": {},
"cost.usd_input": {},
"cost.usd_cached_input": {},
"cost.usd_cache_creation": {},
"cost.usd_output": {},
"cost.usd_total": {},
"cost.usd_cache": {},
"llm.authorising_groups": {},
}
// stripAgentNetworkEntryForUsage returns the entry reduced to what's needed to
// record usage/cost: it drops request detail (host / path / source IP) and any
// prompt capture, keeping the LLM usage metadata plus the caller identity
// (user / auth mechanism) needed for attribution. Shipped when an
// agent-network account has log collection disabled but usage must still be
// collected. logEntry is passed by value, so mutating it here is safe; Metadata
// is replaced with a fresh map rather than mutated in place.
func stripAgentNetworkEntryForUsage(entry logEntry) logEntry {
entry.Host = ""
entry.Path = ""
entry.SourceIP = netip.Addr{}
// Drop the rest of the per-request telemetry too — a usage-only entry
// must carry the LLM usage metadata and caller identity, nothing that
// describes the individual request.
entry.Method = ""
entry.ResponseCode = 0
entry.DurationMs = 0
entry.BytesUpload = 0
entry.BytesDownload = 0
entry.Protocol = ""
if len(entry.Metadata) > 0 {
stripped := make(map[string]string, len(usageMetadataKeys))
for k := range usageMetadataKeys {
if v, ok := entry.Metadata[k]; ok {
stripped[k] = v
}
}
entry.Metadata = stripped
}
return entry
}
func (l *Logger) log(entry logEntry) {
// Fire off the log request in a separate routine.
// This increases the possibility of losing a log message
// (although it should still get logged in the event of an error),
// but it will reduce latency returning the request in the
// middleware.
// There is also a chance that log messages will arrive at
// the server out of order; however, the timestamp should
// allow for resolving that on the server.
now := timestamppb.Now()
select {
case l.logSem <- struct{}{}:
default:
total := l.dropped.Add(1)
l.logger.Debugf("access log send dropped: worker limit reached (total dropped: %d)", total)
return
}
go func() {
defer func() { <-l.logSem }()
logCtx, cancel := context.WithTimeout(context.Background(), logSendTimeout)
defer cancel()
// Only OIDC sessions have a meaningful user identity.
if entry.AuthMechanism != auth.MethodOIDC.String() {
entry.UserID = ""
}
var sourceIP string
if entry.SourceIP.IsValid() {
sourceIP = entry.SourceIP.String()
}
if _, err := l.client.SendAccessLog(logCtx, &proto.SendAccessLogRequest{
Log: &proto.AccessLog{
LogId: entry.ID,
AccountId: string(entry.AccountID),
Timestamp: now,
ServiceId: string(entry.ServiceID),
Host: entry.Host,
Path: entry.Path,
DurationMs: entry.DurationMs,
Method: entry.Method,
ResponseCode: entry.ResponseCode,
SourceIp: sourceIP,
AuthMechanism: entry.AuthMechanism,
UserId: entry.UserID,
AuthSuccess: entry.AuthSuccess,
BytesUpload: entry.BytesUpload,
BytesDownload: entry.BytesDownload,
Protocol: string(entry.Protocol),
Metadata: entry.Metadata,
AgentNetwork: entry.AgentNetwork,
},
}); err != nil {
l.logger.WithFields(log.Fields{
"service_id": entry.ServiceID,
"host": entry.Host,
"path": entry.Path,
"duration": entry.DurationMs,
"method": entry.Method,
"response_code": entry.ResponseCode,
"source_ip": sourceIP,
"auth_mechanism": entry.AuthMechanism,
"user_id": entry.UserID,
"auth_success": entry.AuthSuccess,
"error": err,
}).Error("Error sending access log on gRPC connection")
}
}()
}
// trackUsage records request and byte counts per domain, logging when thresholds are hit.
func (l *Logger) trackUsage(domain string, bytesTransferred int64) {
if domain == "" {
return
}
l.usageMux.Lock()
defer l.usageMux.Unlock()
now := time.Now()
usage, exists := l.domainUsage[domain]
if !exists {
usage = &domainUsage{
requestStartTime: now,
bytesStartTime: now,
lastActivity: now,
}
l.domainUsage[domain] = usage
}
usage.lastActivity = now
usage.requestCount++
if usage.requestCount >= requestThreshold {
elapsed := time.Since(usage.requestStartTime)
l.logger.WithFields(log.Fields{
"domain": domain,
"requests": usage.requestCount,
"duration": elapsed.String(),
}).Infof("domain %s had %d requests over %s", domain, usage.requestCount, elapsed)
usage.requestCount = 0
usage.requestStartTime = now
}
usage.bytesTransferred += bytesTransferred
if usage.bytesTransferred >= bytesThreshold {
elapsed := time.Since(usage.bytesStartTime)
bytesInGB := float64(usage.bytesTransferred) / (1024 * 1024 * 1024)
l.logger.WithFields(log.Fields{
"domain": domain,
"bytes": usage.bytesTransferred,
"bytes_gb": bytesInGB,
"duration": elapsed.String(),
}).Infof("domain %s transferred %.2f GB over %s", domain, bytesInGB, elapsed)
usage.bytesTransferred = 0
usage.bytesStartTime = now
}
}
// cleanupStaleUsage removes usage and deny-rate-limit entries that have been inactive.
func (l *Logger) cleanupStaleUsage(ctx context.Context) {
ticker := time.NewTicker(usageCleanupPeriod)
defer ticker.Stop()
for {
select {
case <-ctx.Done():
return
case <-ticker.C:
now := time.Now()
l.cleanupDomainUsage(now)
l.cleanupDenyBuckets(now)
}
}
}
func (l *Logger) cleanupDomainUsage(now time.Time) {
l.usageMux.Lock()
defer l.usageMux.Unlock()
removed := 0
for domain, usage := range l.domainUsage {
if now.Sub(usage.lastActivity) > usageInactiveWindow {
delete(l.domainUsage, domain)
removed++
}
}
if removed > 0 {
l.logger.Debugf("cleaned up %d stale domain usage entries", removed)
}
}
func (l *Logger) cleanupDenyBuckets(now time.Time) {
l.denyMu.Lock()
defer l.denyMu.Unlock()
removed := 0
for key, bucket := range l.denyBuckets {
if now.Sub(bucket.lastLogged) > usageInactiveWindow {
delete(l.denyBuckets, key)
removed++
}
}
if removed > 0 {
l.logger.Debugf("cleaned up %d stale deny rate-limit entries", removed)
}
}
proxy/internal/accesslog/middleware.go
+107
View File
@@ -0,0 +1,107 @@
package accesslog
import (
"net"
"net/http"
"strings"
"time"
"github.com/rs/xid"
"github.com/netbirdio/netbird/proxy/internal/proxy"
"github.com/netbirdio/netbird/proxy/internal/responsewriter"
"github.com/netbirdio/netbird/proxy/web"
)
// Middleware wraps an HTTP handler to log access entries and resolve client IPs.
func (l *Logger) Middleware(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
// Skip logging for internal proxy assets (CSS, JS, etc.)
if strings.HasPrefix(r.URL.Path, web.PathPrefix+"/") {
next.ServeHTTP(w, r)
return
}
// Generate request ID early so it can be used by error pages and log correlation.
requestID := xid.New().String()
l.logger.Debugf("request: request_id=%s method=%s host=%s path=%s", requestID, r.Method, r.Host, r.URL.Path)
// Use a response writer wrapper so we can access the status code later.
sw := &statusWriter{
PassthroughWriter: responsewriter.New(w),
status: http.StatusOK,
}
var bytesRead int64
if r.Body != nil {
r.Body = &bodyCounter{
ReadCloser: r.Body,
bytesRead: &bytesRead,
}
}
// Resolve the source IP using trusted proxy configuration before passing
// the request on, as the proxy will modify forwarding headers.
sourceIp := extractSourceIP(r, l.trustedProxies)
// Create a mutable struct to capture data from downstream handlers.
// We pass a pointer in the context - the pointer itself flows down immutably,
// but the struct it points to can be mutated by inner handlers.
capturedData := proxy.NewCapturedData(requestID)
capturedData.SetClientIP(sourceIp)
ctx := proxy.WithCapturedData(r.Context(), capturedData)
start := time.Now()
next.ServeHTTP(sw, r.WithContext(ctx))
duration := time.Since(start)
host, _, err := net.SplitHostPort(r.Host)
if err != nil {
// Fallback to just using the full host value.
host = r.Host
}
bytesUpload := bytesRead
bytesDownload := sw.bytesWritten
entry := logEntry{
ID: requestID,
ServiceID: capturedData.GetServiceID(),
AccountID: capturedData.GetAccountID(),
Host: host,
Path: r.URL.Path,
DurationMs: duration.Milliseconds(),
Method: r.Method,
ResponseCode: int32(sw.status),
SourceIP: sourceIp,
AuthMechanism: capturedData.GetAuthMethod(),
UserID: capturedData.GetUserID(),
AuthSuccess: sw.status != http.StatusUnauthorized && sw.status != http.StatusForbidden,
BytesUpload: bytesUpload,
BytesDownload: bytesDownload,
Protocol: ProtocolHTTP,
Metadata: capturedData.GetMetadata(),
AgentNetwork: capturedData.GetAgentNetwork(),
}
l.logger.Debugf("response: request_id=%s method=%s host=%s path=%s status=%d duration=%dms source=%s origin=%s service=%s account=%s",
requestID, r.Method, host, r.URL.Path, sw.status, duration.Milliseconds(), sourceIp, capturedData.GetOrigin(), capturedData.GetServiceID(), capturedData.GetAccountID())
// Emit the access log unless the matched target opted out
// (agent-network synth targets do this when the account's
// EnableLogCollection toggle is off). For agent-network entries we
// still ship a stripped, usage-only record even when suppressed, so
// usage/cost is collected regardless of the log-collection toggle;
// request detail and prompt capture are dropped before sending.
switch {
case !capturedData.GetSuppressAccessLog():
l.log(entry)
case entry.AgentNetwork:
l.log(stripAgentNetworkEntryForUsage(entry))
}
// Track usage for cost monitoring (upload + download) by domain
l.trackUsage(host, bytesUpload+bytesDownload)
})
}
proxy/internal/accesslog/middleware_test.go
+185
View File
@@ -0,0 +1,185 @@
package accesslog
import (
"context"
"net/http"
"net/http/httptest"
"sync"
"sync/atomic"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"google.golang.org/grpc"
"github.com/netbirdio/netbird/proxy/internal/proxy"
"github.com/netbirdio/netbird/shared/management/proto"
)
// recorderClient is a minimal stub for the access-log gRPCClient interface. It
// counts SendAccessLog invocations and signals on every call so tests can
// deterministically wait for the goroutine inside Logger.log without sleeping.
type recorderClient struct {
mu sync.Mutex
calls int64
lastEntry *proto.AccessLog
called chan struct{}
}
func newRecorderClient() *recorderClient {
return &recorderClient{called: make(chan struct{}, 16)}
}
func (r *recorderClient) SendAccessLog(_ context.Context, in *proto.SendAccessLogRequest, _ ...grpc.CallOption) (*proto.SendAccessLogResponse, error) {
r.mu.Lock()
r.calls++
r.lastEntry = in.GetLog()
r.mu.Unlock()
select {
case r.called <- struct{}{}:
default:
}
return &proto.SendAccessLogResponse{}, nil
}
func (r *recorderClient) callCount() int64 {
r.mu.Lock()
defer r.mu.Unlock()
return r.calls
}
// newTestLogger builds a Logger backed by the supplied recorderClient. It is
// the same constructor production uses, just with a stub gRPC client — no
// mocks, no interface re-implementations.
func newTestLogger(t *testing.T, client *recorderClient) *Logger {
t.Helper()
logger := NewLogger(client, nil, nil)
t.Cleanup(logger.Close)
return logger
}
// TestMiddleware_SuppressAccessLog_SkipsLogSink asserts the suppression gate.
// When the inner handler stamps SuppressAccessLog=true on CapturedData (mirrors
// what reverseproxy does when the matched target's DisableAccessLog flag is
// set), the middleware must NOT invoke the access-log sink. Bandwidth telemetry
// (trackUsage) keeps running — it's the call to SendAccessLog that we gate.
func TestMiddleware_SuppressAccessLog_SkipsLogSink(t *testing.T) {
client := newRecorderClient()
l := newTestLogger(t, client)
inner := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
cd := proxy.CapturedDataFromContext(r.Context())
require.NotNil(t, cd, "middleware must inject CapturedData into the request context")
cd.SetSuppressAccessLog(true)
w.WriteHeader(http.StatusOK)
})
srv := httptest.NewServer(l.Middleware(inner))
t.Cleanup(srv.Close)
resp, err := http.Get(srv.URL + "/agent-network/v1/chat/completions")
require.NoError(t, err, "GET against suppressed target must succeed")
require.NoError(t, resp.Body.Close())
require.Equal(t, http.StatusOK, resp.StatusCode, "inner handler must run normally")
// Give the goroutine fence a beat (Logger.log dispatches in a goroutine).
// The negative assertion needs a small window: if a send is going to
// happen, it happens promptly.
select {
case <-client.called:
t.Fatalf("access-log sink must not be invoked when SuppressAccessLog=true (got %d call(s))", client.callCount())
case <-time.After(150 * time.Millisecond):
}
assert.Equal(t, int64(0), client.callCount(),
"SendAccessLog must not be called for suppressed requests")
}
// TestMiddleware_SuppressAccessLog_DefaultEmitsLog is the regression sanity:
// when nothing sets SuppressAccessLog (the universal default for every
// non-agent-network target), the middleware MUST still emit the access-log
// entry. This is the guarantee that wires-through to the EnableLogCollection
// gate without breaking anyone who isn't opted in.
func TestMiddleware_SuppressAccessLog_DefaultEmitsLog(t *testing.T) {
client := newRecorderClient()
l := newTestLogger(t, client)
var innerRan atomic.Bool
inner := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
innerRan.Store(true)
// Intentionally DO NOT touch SuppressAccessLog — mirrors every
// non-agent-network target.
w.WriteHeader(http.StatusOK)
})
srv := httptest.NewServer(l.Middleware(inner))
t.Cleanup(srv.Close)
resp, err := http.Get(srv.URL + "/service/healthz")
require.NoError(t, err, "GET against default target must succeed")
require.NoError(t, resp.Body.Close())
require.Equal(t, http.StatusOK, resp.StatusCode, "inner handler must run normally")
require.True(t, innerRan.Load(), "inner handler must have run")
select {
case <-client.called:
case <-time.After(2 * time.Second):
t.Fatalf("SendAccessLog must be invoked for non-suppressed requests, none observed (calls=%d)", client.callCount())
}
assert.Equal(t, int64(1), client.callCount(),
"non-suppressed request must produce exactly one access-log send")
}
// TestMiddleware_SuppressAccessLog_PreservesUsageTracking proves the gate is
// surgical: with SuppressAccessLog=true the access-log send is skipped, but
// the per-domain usage tracker still records the bytes transferred. This is
// the cost-monitoring guarantee called out in the gate's comment.
func TestMiddleware_SuppressAccessLog_PreservesUsageTracking(t *testing.T) {
client := newRecorderClient()
l := newTestLogger(t, client)
payload := []byte("ok")
inner := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
cd := proxy.CapturedDataFromContext(r.Context())
require.NotNil(t, cd, "middleware must inject CapturedData")
cd.SetSuppressAccessLog(true)
w.WriteHeader(http.StatusOK)
_, _ = w.Write(payload)
})
srv := httptest.NewServer(l.Middleware(inner))
t.Cleanup(srv.Close)
resp, err := http.Get(srv.URL + "/agent-network/v1/chat/completions")
require.NoError(t, err)
require.NoError(t, resp.Body.Close())
// Allow trackUsage to land — it runs synchronously after l.log(entry) is
// (would have been) called.
time.Sleep(50 * time.Millisecond)
l.usageMux.Lock()
usage, present := l.domainUsage[hostNoPort(srv.URL)]
l.usageMux.Unlock()
require.True(t, present, "domain usage must be tracked even when the access-log is suppressed")
assert.Greater(t, usage.bytesTransferred, int64(0), "bytesTransferred must include the response payload")
assert.Equal(t, int64(0), client.callCount(),
"SendAccessLog must remain suppressed across the response write")
}
// hostNoPort extracts the host name from an httptest server URL. The
// middleware strips the port before keying domain usage, so the test mirrors
// that to look the entry up.
func hostNoPort(url string) string {
// httptest URLs are always "http://127.0.0.1:PORT".
const prefix = "http://"
host := url[len(prefix):]
for i := 0; i < len(host); i++ {
if host[i] == ':' || host[i] == '/' {
return host[:i]
}
}
return host
}
proxy/internal/accesslog/requestip.go
+16
View File
@@ -0,0 +1,16 @@
package accesslog
import (
"net/http"
"net/netip"
"github.com/netbirdio/netbird/trustedproxy"
)
// extractSourceIP resolves the real client IP from the request using trusted
// proxy configuration. When trustedProxies is non-empty and the direct
// connection is from a trusted source, it walks X-Forwarded-For right-to-left
// skipping trusted IPs. Otherwise it returns RemoteAddr directly.
func extractSourceIP(r *http.Request, trustedProxies *trustedproxy.List) netip.Addr {
return trustedProxies.ResolveClientIP(r.RemoteAddr, r.Header.Get("X-Forwarded-For"))
}
proxy/internal/accesslog/statuswriter.go
+39
View File
@@ -0,0 +1,39 @@
package accesslog
import (
"io"
"github.com/netbirdio/netbird/proxy/internal/responsewriter"
)
// statusWriter captures the HTTP status code and bytes written from responses.
// It embeds responsewriter.PassthroughWriter which handles all the optional
// interfaces (Hijacker, Flusher, Pusher) automatically.
type statusWriter struct {
*responsewriter.PassthroughWriter
status int
bytesWritten int64
}
func (w *statusWriter) WriteHeader(status int) {
w.status = status
w.PassthroughWriter.WriteHeader(status)
}
func (w *statusWriter) Write(b []byte) (int, error) {
n, err := w.PassthroughWriter.Write(b)
w.bytesWritten += int64(n)
return n, err
}
// bodyCounter wraps an io.ReadCloser and counts bytes read from the request body.
type bodyCounter struct {
io.ReadCloser
bytesRead *int64
}
func (bc *bodyCounter) Read(p []byte) (int, error) {
n, err := bc.ReadCloser.Read(p)
*bc.bytesRead += int64(n)
return n, err
}
proxy/internal/acme/locker.go
+102
View File
@@ -0,0 +1,102 @@
package acme
import (
"context"
"path/filepath"
log "github.com/sirupsen/logrus"
"github.com/netbirdio/netbird/proxy/internal/flock"
"github.com/netbirdio/netbird/proxy/internal/k8s"
)
// certLocker provides distributed mutual exclusion for certificate operations.
// Implementations must be safe for concurrent use from multiple goroutines.
type certLocker interface {
// Lock acquires an exclusive lock for the given domain.
// It blocks until the lock is acquired, the context is cancelled, or an
// unrecoverable error occurs. The returned function releases the lock;
// callers must call it exactly once when the critical section is complete.
Lock(ctx context.Context, domain string) (unlock func(), err error)
}
// CertLockMethod controls how ACME certificate locks are coordinated.
type CertLockMethod string
const (
// CertLockAuto detects the environment and selects k8s-lease if running
// in a Kubernetes pod, otherwise flock.
CertLockAuto CertLockMethod = "auto"
// CertLockFlock uses advisory file locks via flock(2).
CertLockFlock CertLockMethod = "flock"
// CertLockK8sLease uses Kubernetes coordination Leases.
CertLockK8sLease CertLockMethod = "k8s-lease"
)
func newCertLocker(method CertLockMethod, certDir string, logger *log.Logger) certLocker {
if logger == nil {
logger = log.StandardLogger()
}
if method == "" || method == CertLockAuto {
if k8s.InCluster() {
method = CertLockK8sLease
} else {
method = CertLockFlock
}
logger.Infof("auto-detected cert lock method: %s", method)
}
switch method {
case CertLockK8sLease:
locker, err := newK8sLeaseLocker(logger)
if err != nil {
logger.Warnf("create k8s lease locker, falling back to flock: %v", err)
return newFlockLocker(certDir, logger)
}
logger.Infof("using k8s lease locker in namespace %s", locker.client.Namespace())
return locker
default:
logger.Infof("using flock cert locker in %s", certDir)
return newFlockLocker(certDir, logger)
}
}
type flockLocker struct {
certDir string
logger *log.Logger
}
func newFlockLocker(certDir string, logger *log.Logger) *flockLocker {
if logger == nil {
logger = log.StandardLogger()
}
return &flockLocker{certDir: certDir, logger: logger}
}
// Lock acquires an advisory file lock for the given domain.
func (l *flockLocker) Lock(ctx context.Context, domain string) (func(), error) {
lockPath := filepath.Join(l.certDir, domain+".lock")
lockFile, err := flock.Lock(ctx, lockPath)
if err != nil {
return nil, err
}
// nil lockFile means locking is not supported (non-unix).
if lockFile == nil {
return func() { /* no-op: locking unsupported on this platform */ }, nil
}
return func() {
if err := flock.Unlock(lockFile); err != nil {
l.logger.Debugf("release cert lock for domain %q: %v", domain, err)
}
}, nil
}
type noopLocker struct{}
// Lock is a no-op that always succeeds immediately.
func (noopLocker) Lock(context.Context, string) (func(), error) {
return func() { /* no-op: locker disabled */ }, nil
}
proxy/internal/acme/locker_k8s.go
+197
View File
@@ -0,0 +1,197 @@
package acme
import (
"context"
"errors"
"fmt"
"os"
"time"
log "github.com/sirupsen/logrus"
"github.com/netbirdio/netbird/proxy/internal/k8s"
)
const (
// leaseDurationSec is the Kubernetes Lease TTL. If the holder crashes without
// releasing the lock, other replicas must wait this long before taking over.
// This is intentionally generous: in the worst case two replicas may both
// issue an ACME request for the same domain, which is harmless (the CA
// deduplicates and the cache converges).
leaseDurationSec = 300
retryBaseBackoff = 500 * time.Millisecond
retryMaxBackoff = 10 * time.Second
)
type k8sLeaseLocker struct {
client *k8s.LeaseClient
identity string
logger *log.Logger
}
func newK8sLeaseLocker(logger *log.Logger) (*k8sLeaseLocker, error) {
client, err := k8s.NewLeaseClient()
if err != nil {
return nil, fmt.Errorf("create k8s lease client: %w", err)
}
identity, err := os.Hostname()
if err != nil {
return nil, fmt.Errorf("get hostname: %w", err)
}
return &k8sLeaseLocker{
client: client,
identity: identity,
logger: logger,
}, nil
}
// Lock acquires a Kubernetes Lease for the given domain using optimistic
// concurrency. It retries with exponential backoff until the lease is
// acquired or the context is cancelled.
func (l *k8sLeaseLocker) Lock(ctx context.Context, domain string) (func(), error) {
leaseName := k8s.LeaseNameForDomain(domain)
backoff := retryBaseBackoff
for {
acquired, err := l.tryAcquire(ctx, leaseName, domain)
if err != nil {
return nil, fmt.Errorf("acquire lease %s for %q: %w", leaseName, domain, err)
}
if acquired {
l.logger.Debugf("k8s lease %s acquired for domain %q", leaseName, domain)
return l.unlockFunc(leaseName, domain), nil
}
l.logger.Debugf("k8s lease %s held by another replica, retrying in %s", leaseName, backoff)
timer := time.NewTimer(backoff)
select {
case <-ctx.Done():
timer.Stop()
return nil, ctx.Err()
case <-timer.C:
}
backoff *= 2
if backoff > retryMaxBackoff {
backoff = retryMaxBackoff
}
}
}
// tryAcquire attempts to create or take over a Lease. Returns (true, nil)
// on success, (false, nil) if the lease is held and not stale, or an error.
func (l *k8sLeaseLocker) tryAcquire(ctx context.Context, name, domain string) (bool, error) {
existing, err := l.client.Get(ctx, name)
if err != nil {
return false, err
}
now := k8s.MicroTime{Time: time.Now().UTC()}
dur := int32(leaseDurationSec)
if existing == nil {
lease := &k8s.Lease{
Metadata: k8s.LeaseMetadata{
Name: name,
Annotations: map[string]string{
"netbird.io/domain": domain,
},
},
Spec: k8s.LeaseSpec{
HolderIdentity: &l.identity,
LeaseDurationSeconds: &dur,
AcquireTime: &now,
RenewTime: &now,
},
}
if _, err := l.client.Create(ctx, lease); errors.Is(err, k8s.ErrConflict) {
return false, nil
} else if err != nil {
return false, err
}
return true, nil
}
if !l.canTakeover(existing) {
return false, nil
}
existing.Spec.HolderIdentity = &l.identity
existing.Spec.LeaseDurationSeconds = &dur
existing.Spec.AcquireTime = &now
existing.Spec.RenewTime = &now
if _, err := l.client.Update(ctx, existing); errors.Is(err, k8s.ErrConflict) {
return false, nil
} else if err != nil {
return false, err
}
return true, nil
}
// canTakeover returns true if the lease is free (no holder) or stale
// (renewTime + leaseDuration has passed).
func (l *k8sLeaseLocker) canTakeover(lease *k8s.Lease) bool {
holder := lease.Spec.HolderIdentity
if holder == nil || *holder == "" {
return true
}
// We already hold it (e.g. from a previous crashed attempt).
if *holder == l.identity {
return true
}
if lease.Spec.RenewTime == nil || lease.Spec.LeaseDurationSeconds == nil {
return true
}
expiry := lease.Spec.RenewTime.Add(time.Duration(*lease.Spec.LeaseDurationSeconds) * time.Second)
if time.Now().After(expiry) {
l.logger.Infof("k8s lease %s held by %q is stale (expired %s ago), taking over",
lease.Metadata.Name, *holder, time.Since(expiry).Round(time.Second))
return true
}
return false
}
// unlockFunc returns a closure that releases the lease by clearing the holder.
func (l *k8sLeaseLocker) unlockFunc(name, domain string) func() {
return func() {
// Use a fresh context: the parent may already be cancelled.
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
// Re-GET to get current resourceVersion (ours may be stale if
// the lock was held for a long time and something updated it).
current, err := l.client.Get(ctx, name)
if err != nil {
l.logger.Debugf("release k8s lease %s for %q: get: %v", name, domain, err)
return
}
if current == nil {
return
}
// Only clear if we're still the holder.
if current.Spec.HolderIdentity == nil || *current.Spec.HolderIdentity != l.identity {
l.logger.Debugf("k8s lease %s for %q: holder changed to %v, skip release",
name, domain, current.Spec.HolderIdentity)
return
}
empty := ""
current.Spec.HolderIdentity = &empty
current.Spec.AcquireTime = nil
current.Spec.RenewTime = nil
if _, err := l.client.Update(ctx, current); err != nil {
l.logger.Debugf("release k8s lease %s for %q: update: %v", name, domain, err)
}
}
}
proxy/internal/acme/locker_test.go
+65
View File
@@ -0,0 +1,65 @@
package acme
import (
"context"
"os"
"path/filepath"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestFlockLockerRoundTrip(t *testing.T) {
dir := t.TempDir()
locker := newFlockLocker(dir, nil)
unlock, err := locker.Lock(context.Background(), "example.com")
require.NoError(t, err)
require.NotNil(t, unlock)
// Lock file should exist.
assert.FileExists(t, filepath.Join(dir, "example.com.lock"))
unlock()
}
func TestNoopLocker(t *testing.T) {
locker := noopLocker{}
unlock, err := locker.Lock(context.Background(), "example.com")
require.NoError(t, err)
require.NotNil(t, unlock)
unlock()
}
func TestNewCertLockerDefaultsToFlock(t *testing.T) {
dir := t.TempDir()
// t.Setenv registers cleanup to restore the original value.
// os.Unsetenv is needed because the production code uses LookupEnv,
// which distinguishes "empty" from "not set".
t.Setenv("KUBERNETES_SERVICE_HOST", "")
os.Unsetenv("KUBERNETES_SERVICE_HOST")
locker := newCertLocker(CertLockAuto, dir, nil)
_, ok := locker.(*flockLocker)
assert.True(t, ok, "auto without k8s env should select flockLocker")
}
func TestNewCertLockerExplicitFlock(t *testing.T) {
dir := t.TempDir()
locker := newCertLocker(CertLockFlock, dir, nil)
_, ok := locker.(*flockLocker)
assert.True(t, ok, "explicit flock should select flockLocker")
}
func TestNewCertLockerK8sFallsBackToFlock(t *testing.T) {
dir := t.TempDir()
// k8s-lease without SA files should fall back to flock.
locker := newCertLocker(CertLockK8sLease, dir, nil)
_, ok := locker.(*flockLocker)
assert.True(t, ok, "k8s-lease without SA should fall back to flockLocker")
}
proxy/internal/acme/manager.go
+643
View File
@@ -0,0 +1,643 @@
package acme
import (
"context"
"crypto/tls"
"crypto/x509"
"encoding/asn1"
"encoding/base64"
"encoding/binary"
"encoding/pem"
"fmt"
"math/rand/v2"
"net"
"os"
"path/filepath"
"slices"
"strings"
"sync"
"time"
log "github.com/sirupsen/logrus"
"golang.org/x/crypto/acme"
"golang.org/x/crypto/acme/autocert"
"github.com/netbirdio/netbird/proxy/internal/certwatch"
"github.com/netbirdio/netbird/proxy/internal/types"
"github.com/netbirdio/netbird/shared/management/domain"
)
// OID for the SCT list extension (1.3.6.1.4.1.11129.2.4.2)
var oidSCTList = asn1.ObjectIdentifier{1, 3, 6, 1, 4, 1, 11129, 2, 4, 2}
type certificateNotifier interface {
NotifyCertificateIssued(ctx context.Context, accountID types.AccountID, serviceID types.ServiceID, domain string) error
}
type domainState int
const (
domainPending domainState = iota
domainReady
domainFailed
)
type domainInfo struct {
accountID types.AccountID
serviceID types.ServiceID
state domainState
err string
}
type metricsRecorder interface {
RecordCertificateIssuance(duration time.Duration)
}
// wildcardEntry maps a domain suffix (e.g. ".example.com") to a certwatch
// watcher that hot-reloads the corresponding wildcard certificate from disk.
type wildcardEntry struct {
suffix string // e.g. ".example.com"
pattern string // e.g. "*.example.com"
watcher *certwatch.Watcher
}
// ManagerConfig holds the configuration values for the ACME certificate manager.
type ManagerConfig struct {
// CertDir is the directory used for caching ACME certificates.
CertDir string
// ACMEURL is the ACME directory URL (e.g. Let's Encrypt).
ACMEURL string
// EABKID and EABHMACKey are optional External Account Binding credentials
// required by some CAs (e.g. ZeroSSL). EABHMACKey is the base64
// URL-encoded string provided by the CA.
EABKID string
EABHMACKey string
// LockMethod controls the cross-replica coordination strategy.
LockMethod CertLockMethod
// WildcardDir is an optional path to a directory containing wildcard
// certificate pairs (<name>.crt / <name>.key). Wildcard patterns are
// extracted from the certificates' SAN lists. Domains matching a
// wildcard are served from disk; all others go through ACME.
WildcardDir string
}
// Manager wraps autocert.Manager with domain tracking and cross-replica
// coordination via a pluggable locking strategy. The locker prevents
// duplicate ACME requests when multiple replicas share a certificate cache.
type Manager struct {
*autocert.Manager
certDir string
locker certLocker
mu sync.RWMutex
domains map[domain.Domain]*domainInfo
// wildcards holds all loaded wildcard certificates, keyed by suffix.
wildcards []wildcardEntry
certNotifier certificateNotifier
logger *log.Logger
metrics metricsRecorder
}
// NewManager creates a new ACME certificate manager.
func NewManager(cfg ManagerConfig, notifier certificateNotifier, logger *log.Logger, metrics metricsRecorder) (*Manager, error) {
if logger == nil {
logger = log.StandardLogger()
}
mgr := &Manager{
certDir: cfg.CertDir,
locker: newCertLocker(cfg.LockMethod, cfg.CertDir, logger),
domains: make(map[domain.Domain]*domainInfo),
certNotifier: notifier,
logger: logger,
metrics: metrics,
}
if cfg.WildcardDir != "" {
entries, err := loadWildcardDir(cfg.WildcardDir, logger)
if err != nil {
return nil, fmt.Errorf("load wildcard certificates from %q: %w", cfg.WildcardDir, err)
}
mgr.wildcards = entries
}
var eab *acme.ExternalAccountBinding
if cfg.EABKID != "" && cfg.EABHMACKey != "" {
decodedKey, err := base64.RawURLEncoding.DecodeString(cfg.EABHMACKey)
if err != nil {
logger.Errorf("failed to decode EAB HMAC key: %v", err)
} else {
eab = &acme.ExternalAccountBinding{
KID: cfg.EABKID,
Key: decodedKey,
}
logger.Infof("configured External Account Binding with KID: %s", cfg.EABKID)
}
}
mgr.Manager = &autocert.Manager{
Prompt: autocert.AcceptTOS,
HostPolicy: mgr.hostPolicy,
Cache: autocert.DirCache(cfg.CertDir),
ExternalAccountBinding: eab,
Client: &acme.Client{
DirectoryURL: cfg.ACMEURL,
},
}
return mgr, nil
}
// WatchWildcards starts watching all wildcard certificate files for changes.
// It blocks until ctx is cancelled. It is a no-op if no wildcards are loaded.
func (mgr *Manager) WatchWildcards(ctx context.Context) {
if len(mgr.wildcards) == 0 {
return
}
seen := make(map[*certwatch.Watcher]struct{})
var wg sync.WaitGroup
for i := range mgr.wildcards {
w := mgr.wildcards[i].watcher
if _, ok := seen[w]; ok {
continue
}
seen[w] = struct{}{}
wg.Add(1)
go func() {
defer wg.Done()
w.Watch(ctx)
}()
}
wg.Wait()
}
// loadWildcardDir scans dir for .crt files, pairs each with a matching .key
// file, loads them, and extracts wildcard SANs (*.example.com) to build
// the suffix lookup entries.
func loadWildcardDir(dir string, logger *log.Logger) ([]wildcardEntry, error) {
crtFiles, err := filepath.Glob(filepath.Join(dir, "*.crt"))
if err != nil {
return nil, fmt.Errorf("glob certificate files: %w", err)
}
if len(crtFiles) == 0 {
return nil, fmt.Errorf("no .crt files found in %s", dir)
}
var entries []wildcardEntry
for _, crtPath := range crtFiles {
base := strings.TrimSuffix(filepath.Base(crtPath), ".crt")
keyPath := filepath.Join(dir, base+".key")
if _, err := os.Stat(keyPath); err != nil {
logger.Warnf("skipping %s: no matching key file %s", crtPath, keyPath)
continue
}
watcher, err := certwatch.NewWatcher(crtPath, keyPath, logger)
if err != nil {
logger.Warnf("skipping %s: %v", crtPath, err)
continue
}
leaf := watcher.Leaf()
if leaf == nil {
logger.Warnf("skipping %s: no parsed leaf certificate", crtPath)
continue
}
for _, san := range leaf.DNSNames {
suffix, ok := parseWildcard(san)
if !ok {
continue
}
entries = append(entries, wildcardEntry{
suffix: suffix,
pattern: san,
watcher: watcher,
})
logger.Infof("wildcard certificate loaded: %s (from %s)", san, filepath.Base(crtPath))
}
}
if len(entries) == 0 {
return nil, fmt.Errorf("no wildcard SANs (*.example.com) found in certificates in %s", dir)
}
return entries, nil
}
// parseWildcard validates a wildcard domain pattern like "*.example.com"
// and returns the suffix ".example.com" for matching.
func parseWildcard(pattern string) (suffix string, ok bool) {
if !strings.HasPrefix(pattern, "*.") {
return "", false
}
parent := pattern[1:] // ".example.com"
if strings.Count(parent, ".") < 1 {
return "", false
}
return strings.ToLower(parent), true
}
// findWildcardEntry returns the wildcard entry that covers host, or nil.
func (mgr *Manager) findWildcardEntry(host string) *wildcardEntry {
if len(mgr.wildcards) == 0 {
return nil
}
host = strings.ToLower(host)
for i := range mgr.wildcards {
e := &mgr.wildcards[i]
if !strings.HasSuffix(host, e.suffix) {
continue
}
// Single-level match: prefix before suffix must have no dots.
prefix := strings.TrimSuffix(host, e.suffix)
if len(prefix) > 0 && !strings.Contains(prefix, ".") {
return e
}
}
return nil
}
// WildcardPatterns returns the wildcard patterns that are currently loaded.
func (mgr *Manager) WildcardPatterns() []string {
patterns := make([]string, len(mgr.wildcards))
for i, e := range mgr.wildcards {
patterns[i] = e.pattern
}
slices.Sort(patterns)
return patterns
}
func (mgr *Manager) hostPolicy(_ context.Context, host string) error {
if h, _, err := net.SplitHostPort(host); err == nil {
host = h
}
mgr.mu.RLock()
_, exists := mgr.domains[domain.Domain(host)]
mgr.mu.RUnlock()
if !exists {
return fmt.Errorf("unknown domain %q", host)
}
return nil
}
// GetCertificate returns the TLS certificate for the given ClientHello.
// If the requested domain matches a loaded wildcard, the static wildcard
// certificate is returned. Otherwise, the ACME autocert manager handles
// the request.
func (mgr *Manager) GetCertificate(hello *tls.ClientHelloInfo) (*tls.Certificate, error) {
if e := mgr.findWildcardEntry(hello.ServerName); e != nil {
return e.watcher.GetCertificate(hello)
}
return mgr.Manager.GetCertificate(hello)
}
// AddDomain registers a domain for certificate management. Domains that
// match a loaded wildcard are marked ready immediately (they use the
// static wildcard certificate) and the method returns true. All other
// domains go through ACME prefetch and the method returns false.
//
// When AddDomain returns true the caller is responsible for sending any
// certificate-ready notifications after the surrounding operation (e.g.
// mapping update) has committed successfully.
func (mgr *Manager) AddDomain(d domain.Domain, accountID types.AccountID, serviceID types.ServiceID) (wildcardHit bool) {
name := d.PunycodeString()
if e := mgr.findWildcardEntry(name); e != nil {
mgr.mu.Lock()
mgr.domains[d] = &domainInfo{
accountID: accountID,
serviceID: serviceID,
state: domainReady,
}
mgr.mu.Unlock()
mgr.logger.Debugf("domain %q matches wildcard %q, using static certificate", name, e.pattern)
return true
}
mgr.mu.Lock()
mgr.domains[d] = &domainInfo{
accountID: accountID,
serviceID: serviceID,
state: domainPending,
}
mgr.mu.Unlock()
go mgr.prefetchCertificate(d)
return false
}
// prefetchCertificate proactively triggers certificate generation for a domain.
// It acquires a distributed lock to prevent multiple replicas from issuing
// duplicate ACME requests. The second replica will block until the first
// finishes, then find the certificate in the cache.
// ACME and periodic disk reads race; whichever produces a valid certificate
// first wins. This handles cases where locking is unreliable and another
// replica already wrote the cert to the shared cache.
func (mgr *Manager) prefetchCertificate(d domain.Domain) {
time.Sleep(time.Duration(rand.IntN(200)) * time.Millisecond)
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute)
defer cancel()
name := d.PunycodeString()
mgr.logger.Infof("acquiring cert lock for domain %q", name)
lockStart := time.Now()
unlock, err := mgr.locker.Lock(ctx, name)
if err != nil {
mgr.logger.Warnf("acquire cert lock for domain %q, proceeding without lock: %v", name, err)
} else {
mgr.logger.Infof("acquired cert lock for domain %q in %s", name, time.Since(lockStart))
defer unlock()
}
if cert, err := mgr.readCertFromDisk(ctx, name); err == nil {
mgr.logger.Infof("certificate for domain %q already on disk, skipping ACME", name)
mgr.recordAndNotify(ctx, d, name, cert, 0)
return
}
// Run ACME in a goroutine so we can race it against periodic disk reads.
// autocert uses its own internal context and cannot be cancelled externally.
type acmeResult struct {
cert *tls.Certificate
err error
}
acmeCh := make(chan acmeResult, 1)
hello := &tls.ClientHelloInfo{ServerName: name, Conn: &dummyConn{ctx: ctx}}
go func() {
cert, err := mgr.GetCertificate(hello)
acmeCh <- acmeResult{cert, err}
}()
start := time.Now()
diskTicker := time.NewTicker(5 * time.Second)
defer diskTicker.Stop()
for {
select {
case res := <-acmeCh:
elapsed := time.Since(start)
if res.err != nil {
mgr.logger.Warnf("prefetch certificate for domain %q in %s: %v", name, elapsed.String(), res.err)
mgr.setDomainState(d, domainFailed, res.err.Error())
return
}
mgr.recordAndNotify(ctx, d, name, res.cert, elapsed)
return
case <-diskTicker.C:
cert, err := mgr.readCertFromDisk(context.Background(), name)
if err != nil {
continue
}
mgr.logger.Infof("certificate for domain %q appeared on disk after %s", name, time.Since(start).Round(time.Millisecond))
// Drain the ACME goroutine before marking ready — autocert holds
// an internal write lock on certState while ACME is in flight.
go func() {
select {
case <-acmeCh:
default:
}
mgr.recordAndNotify(context.Background(), d, name, cert, 0)
}()
return
case <-ctx.Done():
mgr.logger.Warnf("prefetch certificate for domain %q timed out", name)
mgr.setDomainState(d, domainFailed, ctx.Err().Error())
return
}
}
}
// readCertFromDisk reads and parses a certificate directly from the autocert
// DirCache, bypassing autocert's internal certState mutex. Safe to call
// concurrently with an in-flight ACME request for the same domain.
func (mgr *Manager) readCertFromDisk(ctx context.Context, name string) (*tls.Certificate, error) {
if mgr.Cache == nil {
return nil, fmt.Errorf("no cache configured")
}
data, err := mgr.Cache.Get(ctx, name)
if err != nil {
return nil, err
}
privBlock, certsPEM := pem.Decode(data)
if privBlock == nil || !strings.Contains(privBlock.Type, "PRIVATE") {
return nil, fmt.Errorf("no private key in cache for %q", name)
}
cert, err := tls.X509KeyPair(certsPEM, pem.EncodeToMemory(privBlock))
if err != nil {
return nil, fmt.Errorf("parse cached certificate for %q: %w", name, err)
}
if len(cert.Certificate) > 0 {
leaf, err := x509.ParseCertificate(cert.Certificate[0])
if err != nil {
return nil, fmt.Errorf("parse leaf for %q: %w", name, err)
}
if time.Now().After(leaf.NotAfter) {
return nil, fmt.Errorf("cached certificate for %q expired at %s", name, leaf.NotAfter)
}
cert.Leaf = leaf
}
return &cert, nil
}
// recordAndNotify records metrics, marks the domain ready, logs cert details,
// and notifies the cert notifier.
func (mgr *Manager) recordAndNotify(ctx context.Context, d domain.Domain, name string, cert *tls.Certificate, elapsed time.Duration) {
if elapsed > 0 && mgr.metrics != nil {
mgr.metrics.RecordCertificateIssuance(elapsed)
}
mgr.setDomainState(d, domainReady, "")
now := time.Now()
if cert != nil && cert.Leaf != nil {
leaf := cert.Leaf
mgr.logger.Infof("certificate for domain %q ready in %s: serial=%s SANs=%v notBefore=%s, notAfter=%s, now=%s",
name, elapsed.Round(time.Millisecond),
leaf.SerialNumber.Text(16),
leaf.DNSNames,
leaf.NotBefore.UTC().Format(time.RFC3339),
leaf.NotAfter.UTC().Format(time.RFC3339),
now.UTC().Format(time.RFC3339),
)
mgr.logCertificateDetails(name, leaf, now)
} else {
mgr.logger.Infof("certificate for domain %q ready in %s", name, elapsed.Round(time.Millisecond))
}
mgr.mu.RLock()
info := mgr.domains[d]
mgr.mu.RUnlock()
if info != nil && mgr.certNotifier != nil {
if err := mgr.certNotifier.NotifyCertificateIssued(ctx, info.accountID, info.serviceID, name); err != nil {
mgr.logger.Warnf("notify certificate ready for domain %q: %v", name, err)
}
}
}
func (mgr *Manager) setDomainState(d domain.Domain, state domainState, errMsg string) {
mgr.mu.Lock()
defer mgr.mu.Unlock()
if info, ok := mgr.domains[d]; ok {
info.state = state
info.err = errMsg
}
}
// logCertificateDetails logs certificate validity and SCT timestamps.
func (mgr *Manager) logCertificateDetails(domain string, cert *x509.Certificate, now time.Time) {
if cert.NotBefore.After(now) {
mgr.logger.Warnf("certificate for %q NotBefore is in the future by %v", domain, cert.NotBefore.Sub(now))
}
sctTimestamps := mgr.parseSCTTimestamps(cert)
if len(sctTimestamps) == 0 {
return
}
for i, sctTime := range sctTimestamps {
if sctTime.After(now) {
mgr.logger.Warnf("certificate for %q SCT[%d] timestamp is in the future: %v (by %v)",
domain, i, sctTime.UTC(), sctTime.Sub(now))
} else {
mgr.logger.Debugf("certificate for %q SCT[%d] timestamp: %v (%v in the past)",
domain, i, sctTime.UTC(), now.Sub(sctTime))
}
}
}
// parseSCTTimestamps extracts SCT timestamps from a certificate.
func (mgr *Manager) parseSCTTimestamps(cert *x509.Certificate) []time.Time {
var timestamps []time.Time
for _, ext := range cert.Extensions {
if !ext.Id.Equal(oidSCTList) {
continue
}
// The extension value is an OCTET STRING containing the SCT list
var sctListBytes []byte
if _, err := asn1.Unmarshal(ext.Value, &sctListBytes); err != nil {
mgr.logger.Debugf("failed to unmarshal SCT list outer wrapper: %v", err)
continue
}
// SCT list format: 2-byte length prefix, then concatenated SCTs
if len(sctListBytes) < 2 {
continue
}
listLen := int(binary.BigEndian.Uint16(sctListBytes[:2]))
data := sctListBytes[2:]
if len(data) < listLen {
continue
}
// Parse individual SCTs
offset := 0
for offset < listLen {
if offset+2 > len(data) {
break
}
sctLen := int(binary.BigEndian.Uint16(data[offset : offset+2]))
offset += 2
if offset+sctLen > len(data) {
break
}
sctData := data[offset : offset+sctLen]
offset += sctLen
// SCT format: version (1) + log_id (32) + timestamp (8) + ...
if len(sctData) < 41 {
continue
}
// Timestamp is at offset 33 (after version + log_id), 8 bytes, milliseconds since epoch
tsMillis := binary.BigEndian.Uint64(sctData[33:41])
ts := time.UnixMilli(int64(tsMillis))
timestamps = append(timestamps, ts)
}
}
return timestamps
}
// dummyConn implements net.Conn to provide context for certificate fetching.
type dummyConn struct {
ctx context.Context
}
func (c *dummyConn) Read(b []byte) (n int, err error) { return 0, nil }
func (c *dummyConn) Write(b []byte) (n int, err error) { return len(b), nil }
func (c *dummyConn) Close() error { return nil }
func (c *dummyConn) LocalAddr() net.Addr { return nil }
func (c *dummyConn) RemoteAddr() net.Addr { return nil }
func (c *dummyConn) SetDeadline(t time.Time) error { return nil }
func (c *dummyConn) SetReadDeadline(t time.Time) error { return nil }
func (c *dummyConn) SetWriteDeadline(t time.Time) error { return nil }
// RemoveDomain removes a domain from tracking.
func (mgr *Manager) RemoveDomain(d domain.Domain) {
mgr.mu.Lock()
defer mgr.mu.Unlock()
delete(mgr.domains, d)
}
// PendingCerts returns the number of certificates currently being prefetched.
func (mgr *Manager) PendingCerts() int {
mgr.mu.RLock()
defer mgr.mu.RUnlock()
var n int
for _, info := range mgr.domains {
if info.state == domainPending {
n++
}
}
return n
}
// TotalDomains returns the total number of registered domains.
func (mgr *Manager) TotalDomains() int {
mgr.mu.RLock()
defer mgr.mu.RUnlock()
return len(mgr.domains)
}
// PendingDomains returns the domain names currently being prefetched.
func (mgr *Manager) PendingDomains() []string {
return mgr.domainsByState(domainPending)
}
// ReadyDomains returns domain names that have successfully obtained certificates.
func (mgr *Manager) ReadyDomains() []string {
return mgr.domainsByState(domainReady)
}
// FailedDomains returns domain names that failed certificate prefetch, mapped to their error.
func (mgr *Manager) FailedDomains() map[string]string {
mgr.mu.RLock()
defer mgr.mu.RUnlock()
result := make(map[string]string)
for d, info := range mgr.domains {
if info.state == domainFailed {
result[d.PunycodeString()] = info.err
}
}
return result
}
func (mgr *Manager) domainsByState(state domainState) []string {
mgr.mu.RLock()
defer mgr.mu.RUnlock()
var domains []string
for d, info := range mgr.domains {
if info.state == state {
domains = append(domains, d.PunycodeString())
}
}
slices.Sort(domains)
return domains
}
proxy/internal/acme/manager_test.go
+306
View File
@@ -0,0 +1,306 @@
package acme
import (
"context"
"crypto/ecdsa"
"crypto/elliptic"
"crypto/rand"
"crypto/tls"
"crypto/x509"
"crypto/x509/pkix"
"encoding/pem"
"math/big"
"os"
"path/filepath"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/netbirdio/netbird/proxy/internal/types"
)
func TestHostPolicy(t *testing.T) {
mgr, err := NewManager(ManagerConfig{CertDir: t.TempDir(), ACMEURL: "https://acme.example.com/directory"}, nil, nil, nil)
require.NoError(t, err)
mgr.AddDomain("example.com", types.AccountID("acc1"), types.ServiceID("rp1"))
// Wait for the background prefetch goroutine to finish so the temp dir
// can be cleaned up without a race.
t.Cleanup(func() {
assert.Eventually(t, func() bool {
return mgr.PendingCerts() == 0
}, 30*time.Second, 50*time.Millisecond)
})
tests := []struct {
name string
host string
wantErr bool
}{
{
name: "exact domain match",
host: "example.com",
},
{
name: "domain with port",
host: "example.com:443",
},
{
name: "unknown domain",
host: "unknown.com",
wantErr: true,
},
{
name: "unknown domain with port",
host: "unknown.com:443",
wantErr: true,
},
{
name: "empty host",
host: "",
wantErr: true,
},
{
name: "port only",
host: ":443",
wantErr: true,
},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
err := mgr.hostPolicy(context.Background(), tc.host)
if tc.wantErr {
require.Error(t, err)
assert.Contains(t, err.Error(), "unknown domain")
} else {
assert.NoError(t, err)
}
})
}
}
func TestDomainStates(t *testing.T) {
mgr, err := NewManager(ManagerConfig{CertDir: t.TempDir(), ACMEURL: "https://acme.example.com/directory"}, nil, nil, nil)
require.NoError(t, err)
assert.Equal(t, 0, mgr.PendingCerts(), "initially zero")
assert.Equal(t, 0, mgr.TotalDomains(), "initially zero domains")
assert.Empty(t, mgr.PendingDomains())
assert.Empty(t, mgr.ReadyDomains())
assert.Empty(t, mgr.FailedDomains())
// AddDomain starts as pending, then the prefetch goroutine will fail
// (no real ACME server) and transition to failed.
mgr.AddDomain("a.example.com", types.AccountID("acc1"), types.ServiceID("rp1"))
mgr.AddDomain("b.example.com", types.AccountID("acc1"), types.ServiceID("rp1"))
assert.Equal(t, 2, mgr.TotalDomains(), "two domains registered")
// Pending domains should eventually drain after prefetch goroutines finish.
assert.Eventually(t, func() bool {
return mgr.PendingCerts() == 0
}, 30*time.Second, 100*time.Millisecond, "pending certs should return to zero after prefetch completes")
assert.Empty(t, mgr.PendingDomains())
assert.Equal(t, 2, mgr.TotalDomains(), "total domains unchanged")
// With a fake ACME URL, both should have failed.
failed := mgr.FailedDomains()
assert.Len(t, failed, 2, "both domains should have failed")
assert.Contains(t, failed, "a.example.com")
assert.Contains(t, failed, "b.example.com")
assert.Empty(t, mgr.ReadyDomains())
}
func TestParseWildcard(t *testing.T) {
tests := []struct {
pattern string
wantSuffix string
wantOK bool
}{
{"*.example.com", ".example.com", true},
{"*.foo.example.com", ".foo.example.com", true},
{"*.COM", ".com", true}, // single-label TLD
{"example.com", "", false}, // no wildcard prefix
{"*example.com", "", false}, // missing dot
{"**.example.com", "", false}, // double star
{"", "", false},
}
for _, tc := range tests {
t.Run(tc.pattern, func(t *testing.T) {
suffix, ok := parseWildcard(tc.pattern)
assert.Equal(t, tc.wantOK, ok)
if ok {
assert.Equal(t, tc.wantSuffix, suffix)
}
})
}
}
func TestMatchesWildcard(t *testing.T) {
wcDir := t.TempDir()
generateSelfSignedCert(t, wcDir, "example", "*.example.com")
acmeDir := t.TempDir()
mgr, err := NewManager(ManagerConfig{CertDir: acmeDir, ACMEURL: "https://acme.example.com/directory", WildcardDir: wcDir}, nil, nil, nil)
require.NoError(t, err)
tests := []struct {
host string
match bool
}{
{"foo.example.com", true},
{"bar.example.com", true},
{"FOO.Example.COM", true}, // case insensitive
{"example.com", false}, // bare parent
{"sub.foo.example.com", false}, // multi-level
{"notexample.com", false},
{"", false},
}
for _, tc := range tests {
t.Run(tc.host, func(t *testing.T) {
assert.Equal(t, tc.match, mgr.findWildcardEntry(tc.host) != nil)
})
}
}
// generateSelfSignedCert creates a temporary self-signed certificate and key
// for testing purposes. The baseName controls the output filenames:
// <baseName>.crt and <baseName>.key.
func generateSelfSignedCert(t *testing.T, dir, baseName string, dnsNames ...string) {
t.Helper()
key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
require.NoError(t, err)
template := &x509.Certificate{
SerialNumber: big.NewInt(1),
Subject: pkix.Name{CommonName: dnsNames[0]},
DNSNames: dnsNames,
NotBefore: time.Now().Add(-time.Hour),
NotAfter: time.Now().Add(24 * time.Hour),
}
certDER, err := x509.CreateCertificate(rand.Reader, template, template, &key.PublicKey, key)
require.NoError(t, err)
certFile, err := os.Create(filepath.Join(dir, baseName+".crt"))
require.NoError(t, err)
require.NoError(t, pem.Encode(certFile, &pem.Block{Type: "CERTIFICATE", Bytes: certDER}))
require.NoError(t, certFile.Close())
keyDER, err := x509.MarshalECPrivateKey(key)
require.NoError(t, err)
keyFile, err := os.Create(filepath.Join(dir, baseName+".key"))
require.NoError(t, err)
require.NoError(t, pem.Encode(keyFile, &pem.Block{Type: "EC PRIVATE KEY", Bytes: keyDER}))
require.NoError(t, keyFile.Close())
}
func TestWildcardAddDomainSkipsACME(t *testing.T) {
wcDir := t.TempDir()
generateSelfSignedCert(t, wcDir, "example", "*.example.com")
acmeDir := t.TempDir()
mgr, err := NewManager(ManagerConfig{CertDir: acmeDir, ACMEURL: "https://acme.example.com/directory", WildcardDir: wcDir}, nil, nil, nil)
require.NoError(t, err)
// Add a wildcard-matching domain — should be immediately ready.
mgr.AddDomain("foo.example.com", types.AccountID("acc1"), types.ServiceID("svc1"))
assert.Equal(t, 0, mgr.PendingCerts(), "wildcard domain should not be pending")
assert.Equal(t, []string{"foo.example.com"}, mgr.ReadyDomains())
// Add a non-wildcard domain — should go through ACME (pending then failed).
mgr.AddDomain("other.net", types.AccountID("acc2"), types.ServiceID("svc2"))
assert.Equal(t, 2, mgr.TotalDomains())
// Wait for the ACME prefetch to fail.
assert.Eventually(t, func() bool {
return mgr.PendingCerts() == 0
}, 30*time.Second, 100*time.Millisecond)
assert.Equal(t, []string{"foo.example.com"}, mgr.ReadyDomains())
assert.Contains(t, mgr.FailedDomains(), "other.net")
}
func TestWildcardGetCertificate(t *testing.T) {
wcDir := t.TempDir()
generateSelfSignedCert(t, wcDir, "example", "*.example.com")
acmeDir := t.TempDir()
mgr, err := NewManager(ManagerConfig{CertDir: acmeDir, ACMEURL: "https://acme.example.com/directory", WildcardDir: wcDir}, nil, nil, nil)
require.NoError(t, err)
mgr.AddDomain("foo.example.com", types.AccountID("acc1"), types.ServiceID("svc1"))
// GetCertificate for a wildcard-matching domain should return the static cert.
cert, err := mgr.GetCertificate(&tls.ClientHelloInfo{ServerName: "foo.example.com"})
require.NoError(t, err)
require.NotNil(t, cert)
assert.Contains(t, cert.Leaf.DNSNames, "*.example.com")
}
func TestMultipleWildcards(t *testing.T) {
wcDir := t.TempDir()
generateSelfSignedCert(t, wcDir, "example", "*.example.com")
generateSelfSignedCert(t, wcDir, "other", "*.other.org")
acmeDir := t.TempDir()
mgr, err := NewManager(ManagerConfig{CertDir: acmeDir, ACMEURL: "https://acme.example.com/directory", WildcardDir: wcDir}, nil, nil, nil)
require.NoError(t, err)
assert.ElementsMatch(t, []string{"*.example.com", "*.other.org"}, mgr.WildcardPatterns())
// Both wildcards should resolve.
mgr.AddDomain("foo.example.com", types.AccountID("acc1"), types.ServiceID("svc1"))
mgr.AddDomain("bar.other.org", types.AccountID("acc2"), types.ServiceID("svc2"))
assert.Equal(t, 0, mgr.PendingCerts())
assert.ElementsMatch(t, []string{"foo.example.com", "bar.other.org"}, mgr.ReadyDomains())
// GetCertificate routes to the correct cert.
cert1, err := mgr.GetCertificate(&tls.ClientHelloInfo{ServerName: "foo.example.com"})
require.NoError(t, err)
assert.Contains(t, cert1.Leaf.DNSNames, "*.example.com")
cert2, err := mgr.GetCertificate(&tls.ClientHelloInfo{ServerName: "bar.other.org"})
require.NoError(t, err)
assert.Contains(t, cert2.Leaf.DNSNames, "*.other.org")
// Non-matching domain falls through to ACME.
mgr.AddDomain("custom.net", types.AccountID("acc3"), types.ServiceID("svc3"))
assert.Eventually(t, func() bool {
return mgr.PendingCerts() == 0
}, 30*time.Second, 100*time.Millisecond)
assert.Contains(t, mgr.FailedDomains(), "custom.net")
}
func TestWildcardDirEmpty(t *testing.T) {
wcDir := t.TempDir()
// Empty directory — no .crt files.
_, err := NewManager(ManagerConfig{CertDir: t.TempDir(), ACMEURL: "https://acme.example.com/directory", WildcardDir: wcDir}, nil, nil, nil)
require.Error(t, err)
assert.Contains(t, err.Error(), "no .crt files found")
}
func TestWildcardDirNonWildcardCert(t *testing.T) {
wcDir := t.TempDir()
// Certificate without a wildcard SAN.
generateSelfSignedCert(t, wcDir, "plain", "plain.example.com")
_, err := NewManager(ManagerConfig{CertDir: t.TempDir(), ACMEURL: "https://acme.example.com/directory", WildcardDir: wcDir}, nil, nil, nil)
require.Error(t, err)
assert.Contains(t, err.Error(), "no wildcard SANs")
}
func TestNoWildcardDir(t *testing.T) {
// Empty string means no wildcard dir — pure ACME mode.
mgr, err := NewManager(ManagerConfig{CertDir: t.TempDir(), ACMEURL: "https://acme.example.com/directory"}, nil, nil, nil)
require.NoError(t, err)
assert.Empty(t, mgr.WildcardPatterns())
}
proxy/internal/auth/auth.gohtml
+18
View File
@@ -0,0 +1,18 @@
<!doctype html>
{{ range $method, $value := .Methods }}
{{ if eq $method "pin" }}
<form>
<label for={{ $value }}>PIN:</label>
<input name={{ $value }} id={{ $value }} />
<button type=submit>Submit</button>
</form>
{{ else if eq $method "password" }}
<form>
<label for={{ $value }}>Password:</label>
<input name={{ $value }} id={{ $value }}/>
<button type=submit>Submit</button>
</form>
{{ else if eq $method "oidc" }}
<a href={{ $value }}>Click here to log in with SSO</a>
{{ end }}
{{ end }}
proxy/internal/auth/header.go
+69
View File
@@ -0,0 +1,69 @@
package auth
import (
"errors"
"fmt"
"net/http"
"github.com/netbirdio/netbird/proxy/auth"
"github.com/netbirdio/netbird/proxy/internal/types"
"github.com/netbirdio/netbird/shared/management/proto"
)
// ErrHeaderAuthFailed indicates that the header was present but the
// credential did not validate. Callers should return 401 instead of
// falling through to other auth schemes.
var ErrHeaderAuthFailed = errors.New("header authentication failed")
// Header implements header-based authentication. The proxy checks for the
// configured header in each request and validates its value via gRPC.
type Header struct {
id types.ServiceID
accountId types.AccountID
headerName string
client authenticator
}
// NewHeader creates a Header authentication scheme for the given header name.
func NewHeader(client authenticator, id types.ServiceID, accountId types.AccountID, headerName string) Header {
return Header{
id: id,
accountId: accountId,
headerName: headerName,
client: client,
}
}
// Type returns auth.MethodHeader.
func (Header) Type() auth.Method {
return auth.MethodHeader
}
// Authenticate checks for the configured header in the request. If absent,
// returns empty (unauthenticated). If present, validates via gRPC.
func (h Header) Authenticate(r *http.Request) (string, string, error) {
value := r.Header.Get(h.headerName)
if value == "" {
return "", "", nil
}
res, err := h.client.Authenticate(r.Context(), &proto.AuthenticateRequest{
Id: string(h.id),
AccountId: string(h.accountId),
Request: &proto.AuthenticateRequest_HeaderAuth{
HeaderAuth: &proto.HeaderAuthRequest{
HeaderValue: value,
HeaderName: h.headerName,
},
},
})
if err != nil {
return "", "", fmt.Errorf("authenticate header: %w", err)
}
if res.GetSuccess() {
return res.GetSessionToken(), "", nil
}
return "", "", ErrHeaderAuthFailed
}
proxy/internal/auth/identity.go
+47
View File
@@ -0,0 +1,47 @@
package auth
import (
"context"
"net/netip"
)
// PeerIdentity describes the locally-known facts about a peer reachable on
// the proxy's per-account WireGuard listener. Phase 3 fills PubKey, TunnelIP
// and FQDN from the embedded client's peerstore. UserID, Email and Groups
// stay zero in V1 — full identity still travels through ValidateTunnelPeer.
// Phase V2 will populate them once RemotePeerConfig carries user identity.
type PeerIdentity struct {
PubKey string
TunnelIP netip.Addr
FQDN string
// V2 fields (zero in V1).
UserID string
Email string
Groups []string
}
// TunnelLookupFunc resolves a tunnel IP to a peer identity using locally
// available peerstore data. ok=false means the IP is not in the calling
// account's roster.
type TunnelLookupFunc func(ip netip.Addr) (PeerIdentity, bool)
type tunnelLookupContextKey struct{}
// WithTunnelLookup attaches a per-account peerstore lookup function to
// the request context. The auth middleware calls this lookup before
// hitting management's ValidateTunnelPeer to short-circuit unknown IPs
// and to skip the RPC for already-cached identities.
func WithTunnelLookup(ctx context.Context, lookup TunnelLookupFunc) context.Context {
if lookup == nil {
return ctx
}
return context.WithValue(ctx, tunnelLookupContextKey{}, lookup)
}
// TunnelLookupFromContext returns the peerstore lookup attached to ctx,
// or nil when the request did not arrive on a per-account listener.
func TunnelLookupFromContext(ctx context.Context) TunnelLookupFunc {
v, _ := ctx.Value(tunnelLookupContextKey{}).(TunnelLookupFunc)
return v
}
proxy/internal/auth/middleware.go
+740
View File
@@ -0,0 +1,740 @@
package auth
import (
"context"
"crypto/ed25519"
"encoding/base64"
"errors"
"fmt"
"html"
"net"
"net/http"
"net/netip"
"net/url"
"sync"
"time"
log "github.com/sirupsen/logrus"
"google.golang.org/grpc"
"github.com/netbirdio/netbird/proxy/auth"
"github.com/netbirdio/netbird/proxy/internal/proxy"
"github.com/netbirdio/netbird/proxy/internal/restrict"
"github.com/netbirdio/netbird/proxy/internal/types"
"github.com/netbirdio/netbird/proxy/web"
"github.com/netbirdio/netbird/shared/management/proto"
)
// errValidationUnavailable indicates that session validation failed due to
// an infrastructure error (e.g. gRPC unavailable), not an invalid token.
var errValidationUnavailable = errors.New("session validation unavailable")
type authenticator interface {
Authenticate(ctx context.Context, in *proto.AuthenticateRequest, opts ...grpc.CallOption) (*proto.AuthenticateResponse, error)
}
// SessionValidator validates session tokens and checks user access permissions.
type SessionValidator interface {
ValidateSession(ctx context.Context, in *proto.ValidateSessionRequest, opts ...grpc.CallOption) (*proto.ValidateSessionResponse, error)
ValidateTunnelPeer(ctx context.Context, in *proto.ValidateTunnelPeerRequest, opts ...grpc.CallOption) (*proto.ValidateTunnelPeerResponse, error)
}
// Scheme defines an authentication mechanism for a domain.
type Scheme interface {
Type() auth.Method
// Authenticate checks the request and determines whether it represents
// an authenticated user. An empty token indicates an unauthenticated
// request; optionally, promptData may be returned for the login UI.
// An error indicates an infrastructure failure (e.g. gRPC unavailable).
Authenticate(*http.Request) (token string, promptData string, err error)
}
// DomainConfig holds the authentication and restriction settings for a protected domain.
type DomainConfig struct {
Schemes []Scheme
SessionPublicKey ed25519.PublicKey
SessionExpiration time.Duration
AccountID types.AccountID
ServiceID types.ServiceID
IPRestrictions *restrict.Filter
// Private routes the domain through ValidateTunnelPeer; failure → 403.
Private bool
}
type validationResult struct {
UserID string
UserEmail string
Valid bool
DeniedReason string
Groups []string
// GroupNames carries the human-readable display names for Groups,
// ordered identically (positional pairing). May be shorter than
// Groups for tokens minted before names were embedded; the consumer
// falls back to ids for missing positions.
GroupNames []string
}
// Middleware applies per-domain authentication and IP restriction checks.
type Middleware struct {
domainsMux sync.RWMutex
domains map[string]DomainConfig
logger *log.Logger
sessionValidator SessionValidator
geo restrict.GeoResolver
tunnelCache *tunnelValidationCache
}
// NewMiddleware creates a new authentication middleware. The sessionValidator is
// optional; if nil, OIDC session tokens are validated locally without group access checks.
func NewMiddleware(logger *log.Logger, sessionValidator SessionValidator, geo restrict.GeoResolver) *Middleware {
if logger == nil {
logger = log.StandardLogger()
}
return &Middleware{
domains: make(map[string]DomainConfig),
logger: logger,
sessionValidator: sessionValidator,
geo: geo,
tunnelCache: newTunnelValidationCache(),
}
}
// Protect wraps next with per-domain authentication and IP restriction checks.
// Requests whose Host is not registered pass through unchanged.
func (mw *Middleware) Protect(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
host, _, err := net.SplitHostPort(r.Host)
if err != nil {
host = r.Host
}
config, exists := mw.getDomainConfig(host)
mw.logger.Debugf("checking authentication for host: %s, exists: %t", host, exists)
if !exists {
next.ServeHTTP(w, r)
return
}
// Set account and service IDs in captured data for access logging.
setCapturedIDs(r, config)
if !mw.checkIPRestrictions(w, r, config) {
return
}
// Private services bypass operator schemes and gate on tunnel peer.
if config.Private {
if mw.forwardWithTunnelPeer(w, r, host, config, next) {
return
}
http.Error(w, "Forbidden", http.StatusForbidden)
return
}
// Domains with no authentication schemes pass through after IP checks.
if len(config.Schemes) == 0 {
next.ServeHTTP(w, r)
return
}
if mw.handleOAuthCallbackError(w, r) {
return
}
if mw.forwardWithSessionCookie(w, r, host, config, next) {
return
}
if mw.forwardWithHeaderAuth(w, r, host, config, next) {
return
}
if mw.forwardWithTunnelPeer(w, r, host, config, next) {
return
}
if mw.blockOIDCOnPlainHTTP(w, r, config) {
return
}
mw.authenticateWithSchemes(w, r, host, config)
})
}
// requestIsPlainHTTP reports whether the request arrived without TLS.
// Used to gate cookie-on-plain warnings and the OIDC plain-HTTP block.
func requestIsPlainHTTP(r *http.Request) bool {
return r.TLS == nil
}
// hasOIDCScheme reports whether any of the configured schemes requires
// TLS to round-trip safely with an external IdP.
func hasOIDCScheme(schemes []Scheme) bool {
for _, s := range schemes {
if s.Type() == auth.MethodOIDC {
return true
}
}
return false
}
// blockOIDCOnPlainHTTP fails fast when an OIDC-configured domain is hit
// over plain HTTP. Most IdPs reject http:// redirect URIs, so surfacing
// the misconfiguration here yields a clearer error than the IdP's
// "invalid redirect_uri" round-trip.
func (mw *Middleware) blockOIDCOnPlainHTTP(w http.ResponseWriter, r *http.Request, config DomainConfig) bool {
if !requestIsPlainHTTP(r) {
return false
}
if !hasOIDCScheme(config.Schemes) {
return false
}
mw.logger.WithFields(log.Fields{
"host": r.Host,
"remote": r.RemoteAddr,
}).Warn("OIDC scheme reached on plain HTTP path; rejecting with 400 — use port 443")
http.Error(w, "OIDC requires TLS — use port 443", http.StatusBadRequest)
return true
}
func (mw *Middleware) getDomainConfig(host string) (DomainConfig, bool) {
mw.domainsMux.RLock()
defer mw.domainsMux.RUnlock()
config, exists := mw.domains[host]
return config, exists
}
func setCapturedIDs(r *http.Request, config DomainConfig) {
if cd := proxy.CapturedDataFromContext(r.Context()); cd != nil {
cd.SetAccountID(config.AccountID)
cd.SetServiceID(config.ServiceID)
}
}
// checkIPRestrictions validates the client IP against the domain's IP restrictions.
// Uses the resolved client IP from CapturedData (which accounts for trusted proxies)
// rather than r.RemoteAddr directly.
func (mw *Middleware) checkIPRestrictions(w http.ResponseWriter, r *http.Request, config DomainConfig) bool {
if config.IPRestrictions == nil {
return true
}
clientIP := mw.resolveClientIP(r)
if !clientIP.IsValid() {
mw.logger.Debugf("IP restriction: cannot resolve client address for %q, denying", r.RemoteAddr)
http.Error(w, "Forbidden", http.StatusForbidden)
return false
}
var verdict restrict.Verdict
if types.IsOverlayOrigin(r.Context()) {
// Geo/CrowdSec checks don't apply over the WireGuard overlay:
// the source address is always inside the NetBird CGNAT range,
// which is never in a GeoIP database or a CrowdSec decision
// list. Enforcing them here would either no-op (best case) or
// fail-closed when the geo database is missing.
verdict = config.IPRestrictions.CheckCIDR(clientIP)
} else {
verdict = config.IPRestrictions.Check(clientIP, mw.geo)
}
if verdict == restrict.Allow {
return true
}
if verdict.IsCrowdSec() {
if cd := proxy.CapturedDataFromContext(r.Context()); cd != nil {
cd.SetMetadata("crowdsec_verdict", verdict.String())
if config.IPRestrictions.IsObserveOnly(verdict) {
cd.SetMetadata("crowdsec_mode", "observe")
}
}
}
if config.IPRestrictions.IsObserveOnly(verdict) {
mw.logger.Debugf("CrowdSec observe: would block %s for %s (%s)", clientIP, r.Host, verdict)
return true
}
reason := verdict.String()
mw.blockIPRestriction(r, reason)
http.Error(w, "Forbidden", http.StatusForbidden)
return false
}
// resolveClientIP extracts the real client IP from CapturedData, falling back to r.RemoteAddr.
func (mw *Middleware) resolveClientIP(r *http.Request) netip.Addr {
if cd := proxy.CapturedDataFromContext(r.Context()); cd != nil {
if ip := cd.GetClientIP(); ip.IsValid() {
return ip
}
}
clientIPStr, _, _ := net.SplitHostPort(r.RemoteAddr)
if clientIPStr == "" {
clientIPStr = r.RemoteAddr
}
addr, err := netip.ParseAddr(clientIPStr)
if err != nil {
return netip.Addr{}
}
return addr.Unmap()
}
// blockIPRestriction sets captured data fields for an IP-restriction block event.
func (mw *Middleware) blockIPRestriction(r *http.Request, reason string) {
if cd := proxy.CapturedDataFromContext(r.Context()); cd != nil {
cd.SetOrigin(proxy.OriginAuth)
cd.SetAuthMethod(reason)
}
mw.logger.Debugf("IP restriction: %s for %s", reason, r.RemoteAddr)
}
// handleOAuthCallbackError checks for error query parameters from an OAuth
// callback and renders the access denied page if present.
func (mw *Middleware) handleOAuthCallbackError(w http.ResponseWriter, r *http.Request) bool {
errCode := r.URL.Query().Get("error")
if errCode == "" {
return false
}
var requestID string
if cd := proxy.CapturedDataFromContext(r.Context()); cd != nil {
cd.SetOrigin(proxy.OriginAuth)
cd.SetAuthMethod(auth.MethodOIDC.String())
requestID = cd.GetRequestID()
}
errDesc := r.URL.Query().Get("error_description")
if errDesc == "" {
errDesc = "An error occurred during authentication"
} else {
errDesc = html.EscapeString(errDesc)
}
web.ServeAccessDeniedPage(w, r, http.StatusForbidden, "Access Denied", errDesc, requestID)
return true
}
// forwardWithSessionCookie checks for a valid session cookie and, if found,
// sets the user identity on the request context and forwards to the next handler.
func (mw *Middleware) forwardWithSessionCookie(w http.ResponseWriter, r *http.Request, host string, config DomainConfig, next http.Handler) bool {
cookie, err := r.Cookie(auth.SessionCookieName)
if err != nil {
return false
}
userID, email, method, groups, groupNames, err := auth.ValidateSessionJWT(cookie.Value, host, config.SessionPublicKey)
if err != nil {
return false
}
if cd := proxy.CapturedDataFromContext(r.Context()); cd != nil {
cd.SetUserID(userID)
cd.SetUserEmail(email)
cd.SetUserGroups(groups)
cd.SetUserGroupNames(groupNames)
cd.SetAuthMethod(method)
}
next.ServeHTTP(w, r)
return true
}
// forwardWithTunnelPeer is the OIDC fast-path for requests originating on the
// netbird mesh. When the source IP belongs to a private/CGNAT range the proxy
// asks management to resolve it to a peer/user and to gate by the service's
// distribution_groups. On success the proxy installs the freshly minted JWT
// as a session cookie, sets UserID + Method=oidc on the captured data, and
// forwards directly — operators see the same access-log shape as if the user
// had completed an OIDC redirect. Any failure (private-range mismatch,
// management unreachable, peer unknown, user not in group) returns false so
// the caller falls back to the existing OIDC scheme dispatch.
//
// The fast-path is gated on TunnelLookupFromContext(r.Context()) being
// present — that context value is attached only by the per-account
// inbound (overlay) listener. The host listener never sets it, so a
// public client whose source IP happens to fall inside an RFC1918 / ULA
// / CGNAT range can't impersonate a mesh peer by colliding with a
// tunnel-IP. Once we know the request arrived over WireGuard the
// per-account peerstore lookup is consulted: a miss denies fast (no
// management round-trip), a hit gates the cached ValidateTunnelPeer RPC
// that mints the session JWT.
func (mw *Middleware) forwardWithTunnelPeer(w http.ResponseWriter, r *http.Request, host string, config DomainConfig, next http.Handler) bool {
if mw.sessionValidator == nil {
return false
}
clientIP := mw.resolveClientIP(r)
if !clientIP.IsValid() {
return false
}
// Anti-spoof: only honour the tunnel-peer fast-path on requests that
// were stamped by an overlay listener. Without that marker an
// attacker could send a request from a colliding RFC1918 / CGNAT
// source on the public listener and bypass operator auth.
lookup := TunnelLookupFromContext(r.Context())
if lookup == nil {
return false
}
if !isTunnelSourceIP(clientIP) {
return false
}
if _, ok := lookup(clientIP); !ok {
mw.logger.WithFields(log.Fields{
"host": host,
"remote": clientIP,
}).Debug("local peerstore: tunnel IP not in account roster; denying without RPC")
return false
}
resp, _, err := mw.tunnelCache.fetch(r.Context(), tunnelCacheKey{
accountID: config.AccountID,
tunnelIP: clientIP,
domain: host,
}, mw.validateTunnelPeer)
if err != nil {
mw.logger.WithError(err).Debug("ValidateTunnelPeer failed; falling back to OIDC")
return false
}
if !resp.GetValid() || resp.GetSessionToken() == "" {
return false
}
setSessionCookie(w, resp.GetSessionToken(), config.SessionExpiration)
if cd := proxy.CapturedDataFromContext(r.Context()); cd != nil {
cd.SetOrigin(proxy.OriginAuth)
cd.SetUserID(resp.GetUserId())
cd.SetUserEmail(resp.GetUserEmail())
cd.SetUserGroups(resp.GetPeerGroupIds())
cd.SetUserGroupNames(resp.GetPeerGroupNames())
cd.SetAuthMethod(auth.MethodOIDC.String())
}
next.ServeHTTP(w, r)
return true
}
// validateTunnelPeer adapts the SessionValidator interface to the cache's
// validateTunnelPeerFn signature.
func (mw *Middleware) validateTunnelPeer(ctx context.Context, req *proto.ValidateTunnelPeerRequest) (*proto.ValidateTunnelPeerResponse, error) {
return mw.sessionValidator.ValidateTunnelPeer(ctx, req)
}
// cgnatPrefix covers RFC 6598 100.64.0.0/10, the CGNAT block NetBird
// allocates tunnel addresses from by default. IsPrivate() doesn't include
// it, so we check it explicitly.
var cgnatPrefix = netip.MustParsePrefix("100.64.0.0/10")
// isTunnelSourceIP reports whether ip falls within an address range typical
// of NetBird tunnels: RFC1918 private space, IPv6 ULA, or CGNAT 100.64/10
// (NetBird's default range). Loopback and link-local are excluded — the
// fast-path is meant for peer-to-peer mesh traffic, not localhost.
func isTunnelSourceIP(ip netip.Addr) bool {
if !ip.IsValid() || ip.IsLoopback() || ip.IsLinkLocalUnicast() {
return false
}
if ip.IsPrivate() {
return true
}
return cgnatPrefix.Contains(ip)
}
// forwardWithHeaderAuth checks for a Header auth scheme. If the header validates,
// the request is forwarded directly (no redirect), which is important for API clients.
func (mw *Middleware) forwardWithHeaderAuth(w http.ResponseWriter, r *http.Request, host string, config DomainConfig, next http.Handler) bool {
for _, scheme := range config.Schemes {
hdr, ok := scheme.(Header)
if !ok {
continue
}
handled := mw.tryHeaderScheme(w, r, host, config, hdr, next)
if handled {
return true
}
}
return false
}
func (mw *Middleware) tryHeaderScheme(w http.ResponseWriter, r *http.Request, host string, config DomainConfig, hdr Header, next http.Handler) bool {
token, _, err := hdr.Authenticate(r)
if err != nil {
return mw.handleHeaderAuthError(w, r, err)
}
if token == "" {
return false
}
result, err := mw.validateSessionToken(r.Context(), host, token, config.SessionPublicKey, auth.MethodHeader)
if err != nil {
setHeaderCapturedData(r.Context(), "", "", nil, nil)
status := http.StatusBadRequest
msg := "invalid session token"
if errors.Is(err, errValidationUnavailable) {
status = http.StatusBadGateway
msg = "authentication service unavailable"
}
http.Error(w, msg, status)
return true
}
if !result.Valid {
setHeaderCapturedData(r.Context(), result.UserID, result.UserEmail, result.Groups, result.GroupNames)
http.Error(w, "Unauthorized", http.StatusUnauthorized)
return true
}
setSessionCookie(w, token, config.SessionExpiration)
if cd := proxy.CapturedDataFromContext(r.Context()); cd != nil {
cd.SetUserID(result.UserID)
cd.SetUserEmail(result.UserEmail)
cd.SetUserGroups(result.Groups)
cd.SetUserGroupNames(result.GroupNames)
cd.SetAuthMethod(auth.MethodHeader.String())
}
next.ServeHTTP(w, r)
return true
}
func (mw *Middleware) handleHeaderAuthError(w http.ResponseWriter, r *http.Request, err error) bool {
if errors.Is(err, ErrHeaderAuthFailed) {
setHeaderCapturedData(r.Context(), "", "", nil, nil)
http.Error(w, "Unauthorized", http.StatusUnauthorized)
return true
}
mw.logger.WithField("scheme", "header").Warnf("header auth infrastructure error: %v", err)
if cd := proxy.CapturedDataFromContext(r.Context()); cd != nil {
cd.SetOrigin(proxy.OriginAuth)
}
http.Error(w, "authentication service unavailable", http.StatusBadGateway)
return true
}
func setHeaderCapturedData(ctx context.Context, userID, userEmail string, groups, groupNames []string) {
cd := proxy.CapturedDataFromContext(ctx)
if cd == nil {
return
}
cd.SetOrigin(proxy.OriginAuth)
cd.SetAuthMethod(auth.MethodHeader.String())
cd.SetUserID(userID)
cd.SetUserEmail(userEmail)
cd.SetUserGroups(groups)
cd.SetUserGroupNames(groupNames)
}
// authenticateWithSchemes tries each configured auth scheme in order.
// On success it sets a session cookie and redirects; on failure it renders the login page.
func (mw *Middleware) authenticateWithSchemes(w http.ResponseWriter, r *http.Request, host string, config DomainConfig) {
methods := make(map[string]string)
var attemptedMethod string
for _, scheme := range config.Schemes {
token, promptData, err := scheme.Authenticate(r)
if err != nil {
mw.logger.WithField("scheme", scheme.Type().String()).Warnf("authentication infrastructure error: %v", err)
if cd := proxy.CapturedDataFromContext(r.Context()); cd != nil {
cd.SetOrigin(proxy.OriginAuth)
}
http.Error(w, "authentication service unavailable", http.StatusBadGateway)
return
}
// Track if credentials were submitted but auth failed
if token == "" && wasCredentialSubmitted(r, scheme.Type()) {
attemptedMethod = scheme.Type().String()
}
if token != "" {
mw.handleAuthenticatedToken(w, r, host, token, config, scheme)
return
}
methods[scheme.Type().String()] = promptData
}
if cd := proxy.CapturedDataFromContext(r.Context()); cd != nil {
cd.SetOrigin(proxy.OriginAuth)
if attemptedMethod != "" {
cd.SetAuthMethod(attemptedMethod)
}
}
if oidcURL, ok := methods[auth.MethodOIDC.String()]; ok && len(methods) == 1 && oidcURL != "" {
http.Redirect(w, r, oidcURL, http.StatusFound)
return
}
web.ServeHTTP(w, r, map[string]any{"methods": methods}, http.StatusUnauthorized)
}
// handleAuthenticatedToken validates the token, handles denied access, and on
// success sets a session cookie and redirects to the original URL.
func (mw *Middleware) handleAuthenticatedToken(w http.ResponseWriter, r *http.Request, host, token string, config DomainConfig, scheme Scheme) {
result, err := mw.validateSessionToken(r.Context(), host, token, config.SessionPublicKey, scheme.Type())
if err != nil {
if cd := proxy.CapturedDataFromContext(r.Context()); cd != nil {
cd.SetOrigin(proxy.OriginAuth)
cd.SetAuthMethod(scheme.Type().String())
}
status := http.StatusBadRequest
msg := "invalid session token"
if errors.Is(err, errValidationUnavailable) {
status = http.StatusBadGateway
msg = "authentication service unavailable"
}
http.Error(w, msg, status)
return
}
if !result.Valid {
var requestID string
if cd := proxy.CapturedDataFromContext(r.Context()); cd != nil {
cd.SetOrigin(proxy.OriginAuth)
cd.SetUserID(result.UserID)
cd.SetUserEmail(result.UserEmail)
cd.SetUserGroups(result.Groups)
cd.SetUserGroupNames(result.GroupNames)
cd.SetAuthMethod(scheme.Type().String())
requestID = cd.GetRequestID()
}
web.ServeAccessDeniedPage(w, r, http.StatusForbidden, "Access Denied", "You are not authorized to access this service", requestID)
return
}
setSessionCookie(w, token, config.SessionExpiration)
// Redirect instead of forwarding the auth POST to the backend.
// The browser will follow with a GET carrying the new session cookie.
if cd := proxy.CapturedDataFromContext(r.Context()); cd != nil {
cd.SetOrigin(proxy.OriginAuth)
cd.SetUserID(result.UserID)
cd.SetUserEmail(result.UserEmail)
cd.SetUserGroups(result.Groups)
cd.SetUserGroupNames(result.GroupNames)
cd.SetAuthMethod(scheme.Type().String())
}
redirectURL := stripSessionTokenParam(r.URL)
http.Redirect(w, r, redirectURL, http.StatusSeeOther)
}
// setSessionCookie writes a session cookie with secure defaults.
func setSessionCookie(w http.ResponseWriter, token string, expiration time.Duration) {
if expiration == 0 {
expiration = auth.DefaultSessionExpiry
}
http.SetCookie(w, &http.Cookie{
Name: auth.SessionCookieName,
Value: token,
Path: "/",
HttpOnly: true,
Secure: true,
SameSite: http.SameSiteLaxMode,
MaxAge: int(expiration.Seconds()),
})
}
// wasCredentialSubmitted checks if credentials were submitted for the given auth method.
func wasCredentialSubmitted(r *http.Request, method auth.Method) bool {
switch method {
case auth.MethodPIN:
return r.FormValue("pin") != ""
case auth.MethodPassword:
return r.FormValue("password") != ""
case auth.MethodOIDC:
return r.URL.Query().Get("session_token") != ""
}
return false
}
// AddDomain registers authentication schemes for the given domain. With schemes a valid session public key is required.
// private=true forces ValidateTunnelPeer enforcement (403 on failure) regardless of the schemes list.
func (mw *Middleware) AddDomain(domain string, schemes []Scheme, publicKeyB64 string, expiration time.Duration, accountID types.AccountID, serviceID types.ServiceID, ipRestrictions *restrict.Filter, private bool) error {
if len(schemes) == 0 {
mw.domainsMux.Lock()
defer mw.domainsMux.Unlock()
mw.domains[domain] = DomainConfig{
AccountID: accountID,
ServiceID: serviceID,
IPRestrictions: ipRestrictions,
Private: private,
}
return nil
}
pubKeyBytes, err := base64.StdEncoding.DecodeString(publicKeyB64)
if err != nil {
return fmt.Errorf("decode session public key for domain %s: %w", domain, err)
}
if len(pubKeyBytes) != ed25519.PublicKeySize {
return fmt.Errorf("invalid session public key size for domain %s: got %d, want %d", domain, len(pubKeyBytes), ed25519.PublicKeySize)
}
mw.domainsMux.Lock()
defer mw.domainsMux.Unlock()
mw.domains[domain] = DomainConfig{
Schemes: schemes,
SessionPublicKey: pubKeyBytes,
SessionExpiration: expiration,
AccountID: accountID,
ServiceID: serviceID,
IPRestrictions: ipRestrictions,
Private: private,
}
return nil
}
// RemoveDomain unregisters authentication for the given domain.
func (mw *Middleware) RemoveDomain(domain string) {
mw.domainsMux.Lock()
defer mw.domainsMux.Unlock()
delete(mw.domains, domain)
}
// validateSessionToken validates a session token. OIDC tokens with a configured
// validator go through gRPC for group access checks; other methods validate locally.
func (mw *Middleware) validateSessionToken(ctx context.Context, host, token string, publicKey ed25519.PublicKey, method auth.Method) (*validationResult, error) {
if method == auth.MethodOIDC && mw.sessionValidator != nil {
resp, err := mw.sessionValidator.ValidateSession(ctx, &proto.ValidateSessionRequest{
Domain: host,
SessionToken: token,
})
if err != nil {
return nil, fmt.Errorf("%w: %w", errValidationUnavailable, err)
}
if !resp.Valid {
mw.logger.WithFields(log.Fields{
"domain": host,
"denied_reason": resp.DeniedReason,
"user_id": resp.UserId,
}).Debug("Session validation denied")
return &validationResult{
UserID: resp.UserId,
UserEmail: resp.GetUserEmail(),
Valid: false,
DeniedReason: resp.DeniedReason,
}, nil
}
return &validationResult{
UserID: resp.UserId,
UserEmail: resp.GetUserEmail(),
Valid: true,
Groups: resp.GetPeerGroupIds(),
GroupNames: resp.GetPeerGroupNames(),
}, nil
}
userID, email, _, groups, groupNames, err := auth.ValidateSessionJWT(token, host, publicKey)
if err != nil {
return nil, err
}
return &validationResult{UserID: userID, UserEmail: email, Valid: true, Groups: groups, GroupNames: groupNames}, nil
}
// stripSessionTokenParam returns the request URI with the session_token query
// parameter removed so it doesn't linger in the browser's address bar or history.
func stripSessionTokenParam(u *url.URL) string {
q := u.Query()
if !q.Has("session_token") {
return u.RequestURI()
}
q.Del("session_token")
clean := *u
clean.RawQuery = q.Encode()
return clean.RequestURI()
}
proxy/internal/auth/middleware_test.go
+1406
View File
File diff suppressed because it is too large Load Diff
proxy/internal/auth/oidc.go
+66
View File
@@ -0,0 +1,66 @@
package auth
import (
"context"
"fmt"
"net/http"
"net/url"
"google.golang.org/grpc"
"github.com/netbirdio/netbird/proxy/auth"
"github.com/netbirdio/netbird/proxy/internal/types"
"github.com/netbirdio/netbird/shared/management/proto"
)
type urlGenerator interface {
GetOIDCURL(context.Context, *proto.GetOIDCURLRequest, ...grpc.CallOption) (*proto.GetOIDCURLResponse, error)
}
type OIDC struct {
id types.ServiceID
accountId types.AccountID
forwardedProto string
client urlGenerator
}
// NewOIDC creates a new OIDC authentication scheme
func NewOIDC(client urlGenerator, id types.ServiceID, accountId types.AccountID, forwardedProto string) OIDC {
return OIDC{
id: id,
accountId: accountId,
forwardedProto: forwardedProto,
client: client,
}
}
func (OIDC) Type() auth.Method {
return auth.MethodOIDC
}
// Authenticate checks for an OIDC session token or obtains the OIDC redirect URL.
func (o OIDC) Authenticate(r *http.Request) (string, string, error) {
// Check for the session_token query param (from OIDC redirects).
// The management server passes the token in the URL because it cannot set
// cookies for the proxy's domain (cookies are domain-scoped per RFC 6265).
if token := r.URL.Query().Get("session_token"); token != "" {
return token, "", nil
}
redirectURL := &url.URL{
Scheme: auth.ResolveProto(o.forwardedProto, r.TLS),
Host: r.Host,
Path: r.URL.Path,
}
res, err := o.client.GetOIDCURL(r.Context(), &proto.GetOIDCURLRequest{
Id: string(o.id),
AccountId: string(o.accountId),
RedirectUrl: redirectURL.String(),
})
if err != nil {
return "", "", fmt.Errorf("get OIDC URL: %w", err)
}
return "", res.GetUrl(), nil
}
proxy/internal/auth/password.go
+63
View File
@@ -0,0 +1,63 @@
package auth
import (
"fmt"
"net/http"
"github.com/netbirdio/netbird/proxy/auth"
"github.com/netbirdio/netbird/proxy/internal/types"
"github.com/netbirdio/netbird/shared/management/proto"
)
const passwordFormId = "password"
type Password struct {
id types.ServiceID
accountId types.AccountID
client authenticator
}
func NewPassword(client authenticator, id types.ServiceID, accountId types.AccountID) Password {
return Password{
id: id,
accountId: accountId,
client: client,
}
}
func (Password) Type() auth.Method {
return auth.MethodPassword
}
// Authenticate attempts to authenticate the request using a form
// value passed in the request.
// If authentication fails, the required HTTP form ID is returned
// so that it can be injected into a request from the UI so that
// authentication may be successful.
func (p Password) Authenticate(r *http.Request) (string, string, error) {
password := r.FormValue(passwordFormId)
if password == "" {
// No password submitted; return the form ID so the UI can prompt the user.
return "", passwordFormId, nil
}
res, err := p.client.Authenticate(r.Context(), &proto.AuthenticateRequest{
Id: string(p.id),
AccountId: string(p.accountId),
Request: &proto.AuthenticateRequest_Password{
Password: &proto.PasswordRequest{
Password: password,
},
},
})
if err != nil {
return "", "", fmt.Errorf("authenticate password: %w", err)
}
if res.GetSuccess() {
return res.GetSessionToken(), "", nil
}
return "", passwordFormId, nil
}
proxy/internal/auth/pin.go
+63
View File
@@ -0,0 +1,63 @@
package auth
import (
"fmt"
"net/http"
"github.com/netbirdio/netbird/proxy/auth"
"github.com/netbirdio/netbird/proxy/internal/types"
"github.com/netbirdio/netbird/shared/management/proto"
)
const pinFormId = "pin"
type Pin struct {
id types.ServiceID
accountId types.AccountID
client authenticator
}
func NewPin(client authenticator, id types.ServiceID, accountId types.AccountID) Pin {
return Pin{
id: id,
accountId: accountId,
client: client,
}
}
func (Pin) Type() auth.Method {
return auth.MethodPIN
}
// Authenticate attempts to authenticate the request using a form
// value passed in the request.
// If authentication fails, the required HTTP form ID is returned
// so that it can be injected into a request from the UI so that
// authentication may be successful.
func (p Pin) Authenticate(r *http.Request) (string, string, error) {
pin := r.FormValue(pinFormId)
if pin == "" {
// No PIN submitted; return the form ID so the UI can prompt the user.
return "", pinFormId, nil
}
res, err := p.client.Authenticate(r.Context(), &proto.AuthenticateRequest{
Id: string(p.id),
AccountId: string(p.accountId),
Request: &proto.AuthenticateRequest_Pin{
Pin: &proto.PinRequest{
Pin: pin,
},
},
})
if err != nil {
return "", "", fmt.Errorf("authenticate pin: %w", err)
}
if res.GetSuccess() {
return res.GetSessionToken(), "", nil
}
return "", pinFormId, nil
}
proxy/internal/auth/tunnel_cache.go
+171
View File
@@ -0,0 +1,171 @@
package auth
import (
"context"
"net/netip"
"sync"
"time"
"golang.org/x/sync/singleflight"
"github.com/netbirdio/netbird/proxy/internal/types"
"github.com/netbirdio/netbird/shared/management/proto"
)
// tunnelCacheTTL caps how long a positive ValidateTunnelPeer result is
// reused before re-fetching from management. 5 minutes balances freshness
// against management load on busy mesh networks.
const tunnelCacheTTL = 300 * time.Second
// tunnelCachePerAccount caps the number of cached identities per account.
// Bounded eviction avoids memory growth in pathological cases (huge peer
// roster, brief request bursts) while staying generous for normal use.
const tunnelCachePerAccount = 1024
// tunnelCacheKey identifies a cached entry by tunnel IP and originating
// account. Domain is part of the value, not the key, because the
// management response is per (account, IP) — domain only gates whether a
// re-fetch is needed if the operator is accessing a different service.
type tunnelCacheKey struct {
accountID types.AccountID
tunnelIP netip.Addr
domain string
}
// tunnelCacheEntry stores a positive validation response with the time it
// was minted. Entries past tunnelCacheTTL are treated as misses.
type tunnelCacheEntry struct {
resp *proto.ValidateTunnelPeerResponse
cachedAt time.Time
}
// tunnelValidationCache memoizes ValidateTunnelPeer responses keyed by
// (accountID, tunnelIP, domain). Only successful, valid responses are
// cached — denials skip the cache so policy changes apply immediately.
// Single-flight de-duplicates concurrent fetches for the same key so a
// burst of cold requests collapses into a single RPC.
type tunnelValidationCache struct {
mu sync.Mutex
entries map[types.AccountID]*accountBucket
flight singleflight.Group
ttl time.Duration
maxSize int
now func() time.Time
}
// accountBucket holds the cached entries for a single account, with a
// FIFO eviction queue used when the bucket exceeds maxSize.
type accountBucket struct {
items map[tunnelCacheKey]tunnelCacheEntry
order []tunnelCacheKey
}
// newTunnelValidationCache constructs a cache with default TTL and bounds.
func newTunnelValidationCache() *tunnelValidationCache {
return &tunnelValidationCache{
entries: make(map[types.AccountID]*accountBucket),
ttl: tunnelCacheTTL,
maxSize: tunnelCachePerAccount,
now: time.Now,
}
}
// get returns a cached response for the key, or nil when missing or
// expired. Expired entries are evicted lazily on read.
func (c *tunnelValidationCache) get(key tunnelCacheKey) *proto.ValidateTunnelPeerResponse {
c.mu.Lock()
defer c.mu.Unlock()
bucket, ok := c.entries[key.accountID]
if !ok {
return nil
}
entry, ok := bucket.items[key]
if !ok {
return nil
}
if c.now().Sub(entry.cachedAt) > c.ttl {
delete(bucket.items, key)
bucket.order = removeKey(bucket.order, key)
return nil
}
return entry.resp
}
// put records a positive response under the key. Evicts the oldest entry
// in the account's bucket when the bound is exceeded.
func (c *tunnelValidationCache) put(key tunnelCacheKey, resp *proto.ValidateTunnelPeerResponse) {
c.mu.Lock()
defer c.mu.Unlock()
bucket, ok := c.entries[key.accountID]
if !ok {
bucket = &accountBucket{items: make(map[tunnelCacheKey]tunnelCacheEntry)}
c.entries[key.accountID] = bucket
}
if _, exists := bucket.items[key]; !exists {
bucket.order = append(bucket.order, key)
}
bucket.items[key] = tunnelCacheEntry{resp: resp, cachedAt: c.now()}
for len(bucket.order) > c.maxSize {
oldest := bucket.order[0]
bucket.order = bucket.order[1:]
delete(bucket.items, oldest)
}
}
// removeKey drops the first occurrence of needle from order. The cache
// uses small slices so a linear scan is cheaper than a map+slice combo.
func removeKey(order []tunnelCacheKey, needle tunnelCacheKey) []tunnelCacheKey {
for i, k := range order {
if k == needle {
return append(order[:i], order[i+1:]...)
}
}
return order
}
// flightKey turns a cache key into a single-flight string. AccountID and
// IP isolation by themselves are insufficient because different domains
// for the same peer/account may have different group access.
func flightKey(key tunnelCacheKey) string {
return string(key.accountID) + "|" + key.tunnelIP.String() + "|" + key.domain
}
// validateTunnelPeerFn is the RPC entry point the cache wraps. It matches
// the SessionValidator.ValidateTunnelPeer signature without exposing the
// gRPC option variadic, since callers don't need it on the cache hot path.
type validateTunnelPeerFn func(ctx context.Context, req *proto.ValidateTunnelPeerRequest) (*proto.ValidateTunnelPeerResponse, error)
// fetch returns a cached response when present, otherwise calls validate
// under single-flight and caches the result. Denied responses pass
// through but are not cached so policy changes apply immediately.
func (c *tunnelValidationCache) fetch(ctx context.Context, key tunnelCacheKey, validate validateTunnelPeerFn) (*proto.ValidateTunnelPeerResponse, bool, error) {
if resp := c.get(key); resp != nil {
return resp, true, nil
}
flight := flightKey(key)
res, err, _ := c.flight.Do(flight, func() (any, error) {
if cached := c.get(key); cached != nil {
return cached, nil
}
resp, err := validate(ctx, &proto.ValidateTunnelPeerRequest{
TunnelIp: key.tunnelIP.String(),
Domain: key.domain,
})
if err != nil {
return nil, err
}
if resp.GetValid() && resp.GetSessionToken() != "" {
c.put(key, resp)
}
return resp, nil
})
if err != nil {
return nil, false, err
}
resp, _ := res.(*proto.ValidateTunnelPeerResponse)
return resp, false, nil
}
proxy/internal/auth/tunnel_cache_test.go
+171
View File
@@ -0,0 +1,171 @@
package auth
import (
"context"
"net/netip"
"sync"
"sync/atomic"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/netbirdio/netbird/proxy/internal/types"
"github.com/netbirdio/netbird/shared/management/proto"
)
func newTestKey(account types.AccountID, ip string, domain string) tunnelCacheKey {
return tunnelCacheKey{
accountID: account,
tunnelIP: netip.MustParseAddr(ip),
domain: domain,
}
}
func TestTunnelCache_HitSkipsRPC(t *testing.T) {
cache := newTunnelValidationCache()
key := newTestKey("acct-1", "100.64.0.10", "svc.example")
var calls int32
validate := func(_ context.Context, req *proto.ValidateTunnelPeerRequest) (*proto.ValidateTunnelPeerResponse, error) {
atomic.AddInt32(&calls, 1)
return &proto.ValidateTunnelPeerResponse{Valid: true, SessionToken: "tok", UserId: "user-1"}, nil
}
resp, fromCache, err := cache.fetch(context.Background(), key, validate)
require.NoError(t, err)
require.NotNil(t, resp, "first fetch returns RPC response")
assert.False(t, fromCache, "first fetch must not be cached")
resp2, fromCache2, err := cache.fetch(context.Background(), key, validate)
require.NoError(t, err)
require.NotNil(t, resp2, "second fetch returns cached response")
assert.True(t, fromCache2, "second fetch must be served from cache")
assert.Equal(t, "user-1", resp2.GetUserId(), "cached response should preserve user identity")
assert.Equal(t, int32(1), atomic.LoadInt32(&calls), "validate should run exactly once with one cache hit")
}
func TestTunnelCache_ExpiredEntryRefetches(t *testing.T) {
cache := newTunnelValidationCache()
clock := time.Now()
cache.now = func() time.Time { return clock }
key := newTestKey("acct-1", "100.64.0.10", "svc.example")
var calls int32
validate := func(_ context.Context, _ *proto.ValidateTunnelPeerRequest) (*proto.ValidateTunnelPeerResponse, error) {
atomic.AddInt32(&calls, 1)
return &proto.ValidateTunnelPeerResponse{Valid: true, SessionToken: "tok"}, nil
}
_, _, err := cache.fetch(context.Background(), key, validate)
require.NoError(t, err)
assert.Equal(t, int32(1), atomic.LoadInt32(&calls), "first fetch issues one RPC")
clock = clock.Add(tunnelCacheTTL + time.Second)
_, fromCache, err := cache.fetch(context.Background(), key, validate)
require.NoError(t, err)
assert.False(t, fromCache, "expired entry must miss the cache")
assert.Equal(t, int32(2), atomic.LoadInt32(&calls), "expired entry forces a re-fetch")
}
func TestTunnelCache_DeniedResponseNotCached(t *testing.T) {
cache := newTunnelValidationCache()
key := newTestKey("acct-1", "100.64.0.10", "svc.example")
var calls int32
validate := func(_ context.Context, _ *proto.ValidateTunnelPeerRequest) (*proto.ValidateTunnelPeerResponse, error) {
atomic.AddInt32(&calls, 1)
return &proto.ValidateTunnelPeerResponse{Valid: false, DeniedReason: "not_in_group"}, nil
}
for i := 0; i < 3; i++ {
_, _, err := cache.fetch(context.Background(), key, validate)
require.NoError(t, err, "fetch must not error on denied response")
}
assert.Equal(t, int32(3), atomic.LoadInt32(&calls), "denied responses bypass the cache so policy changes apply immediately")
}
func TestTunnelCache_ConcurrentColdHitsCoalesce(t *testing.T) {
cache := newTunnelValidationCache()
key := newTestKey("acct-1", "100.64.0.10", "svc.example")
gate := make(chan struct{})
var calls int32
validate := func(_ context.Context, _ *proto.ValidateTunnelPeerRequest) (*proto.ValidateTunnelPeerResponse, error) {
atomic.AddInt32(&calls, 1)
<-gate
return &proto.ValidateTunnelPeerResponse{Valid: true, SessionToken: "tok"}, nil
}
const workers = 16
var wg sync.WaitGroup
wg.Add(workers)
results := make([]bool, workers)
for i := 0; i < workers; i++ {
go func(idx int) {
defer wg.Done()
resp, _, err := cache.fetch(context.Background(), key, validate)
results[idx] = err == nil && resp.GetValid()
}(i)
}
time.Sleep(20 * time.Millisecond)
close(gate)
wg.Wait()
for i, ok := range results {
assert.Truef(t, ok, "worker %d should observe a successful response", i)
}
assert.Equal(t, int32(1), atomic.LoadInt32(&calls), "single-flight must collapse concurrent cold fetches into one RPC")
}
func TestTunnelCache_PerAccountIsolation(t *testing.T) {
cache := newTunnelValidationCache()
keyA := newTestKey("acct-a", "100.64.0.10", "svc.example")
keyB := newTestKey("acct-b", "100.64.0.10", "svc.example")
var callsA, callsB int32
validateA := func(_ context.Context, _ *proto.ValidateTunnelPeerRequest) (*proto.ValidateTunnelPeerResponse, error) {
atomic.AddInt32(&callsA, 1)
return &proto.ValidateTunnelPeerResponse{Valid: true, SessionToken: "tok-a", UserId: "user-a"}, nil
}
validateB := func(_ context.Context, _ *proto.ValidateTunnelPeerRequest) (*proto.ValidateTunnelPeerResponse, error) {
atomic.AddInt32(&callsB, 1)
return &proto.ValidateTunnelPeerResponse{Valid: true, SessionToken: "tok-b", UserId: "user-b"}, nil
}
respA, _, err := cache.fetch(context.Background(), keyA, validateA)
require.NoError(t, err)
respB, _, err := cache.fetch(context.Background(), keyB, validateB)
require.NoError(t, err)
assert.Equal(t, "user-a", respA.GetUserId(), "account A response should belong to user-a")
assert.Equal(t, "user-b", respB.GetUserId(), "account B response must not be served from account A's cache")
assert.Equal(t, int32(1), atomic.LoadInt32(&callsA), "validateA called exactly once")
assert.Equal(t, int32(1), atomic.LoadInt32(&callsB), "validateB called exactly once")
}
func TestTunnelCache_BoundedSizeEvictsOldest(t *testing.T) {
cache := newTunnelValidationCache()
cache.maxSize = 2
validate := func(_ context.Context, req *proto.ValidateTunnelPeerRequest) (*proto.ValidateTunnelPeerResponse, error) {
return &proto.ValidateTunnelPeerResponse{Valid: true, SessionToken: "tok-" + req.GetTunnelIp()}, nil
}
keys := []tunnelCacheKey{
newTestKey("acct-1", "100.64.0.10", "svc"),
newTestKey("acct-1", "100.64.0.11", "svc"),
newTestKey("acct-1", "100.64.0.12", "svc"),
}
for _, k := range keys {
_, _, err := cache.fetch(context.Background(), k, validate)
require.NoError(t, err)
}
assert.Nil(t, cache.get(keys[0]), "oldest key should be evicted past maxSize")
assert.NotNil(t, cache.get(keys[1]), "second-newest must remain cached")
assert.NotNil(t, cache.get(keys[2]), "newest must remain cached")
}
proxy/internal/auth/tunnel_lookup_test.go
+355
View File
@@ -0,0 +1,355 @@
package auth
import (
"context"
"errors"
"net/http"
"net/http/httptest"
"net/netip"
"sync/atomic"
"testing"
log "github.com/sirupsen/logrus"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"google.golang.org/grpc"
"github.com/netbirdio/netbird/proxy/internal/proxy"
"github.com/netbirdio/netbird/shared/management/proto"
)
// stubSessionValidator records ValidateTunnelPeer calls and returns the
// pre-canned response. Counts let tests assert RPC traffic.
type stubSessionValidator struct {
respFn func(req *proto.ValidateTunnelPeerRequest) *proto.ValidateTunnelPeerResponse
respErr error
tunnelCalls atomic.Int32
}
func (s *stubSessionValidator) ValidateSession(_ context.Context, _ *proto.ValidateSessionRequest, _ ...grpc.CallOption) (*proto.ValidateSessionResponse, error) {
return &proto.ValidateSessionResponse{Valid: false}, nil
}
func (s *stubSessionValidator) ValidateTunnelPeer(_ context.Context, in *proto.ValidateTunnelPeerRequest, _ ...grpc.CallOption) (*proto.ValidateTunnelPeerResponse, error) {
s.tunnelCalls.Add(1)
if s.respErr != nil {
return nil, s.respErr
}
if s.respFn != nil {
return s.respFn(in), nil
}
return &proto.ValidateTunnelPeerResponse{Valid: false}, nil
}
func newTunnelMiddleware(t *testing.T, validator SessionValidator) *Middleware {
t.Helper()
mw := NewMiddleware(log.New(), validator, nil)
require.NoError(t, mw.AddDomain("svc.example", nil, "", 0, "acct-1", "svc-1", nil, false))
return mw
}
func newTunnelRequest(remoteAddr string) (*httptest.ResponseRecorder, *http.Request) {
w := httptest.NewRecorder()
r := httptest.NewRequest(http.MethodGet, "https://svc.example/", nil)
r.Host = "svc.example"
r.RemoteAddr = remoteAddr
return w, r
}
// TestForwardWithTunnelPeer_LocalLookupUnknownIPDeniesFast verifies the
// short-circuit: a tunnel IP not in the account's roster never reaches
// management's ValidateTunnelPeer.
func TestForwardWithTunnelPeer_LocalLookupUnknownIPDeniesFast(t *testing.T) {
validator := &stubSessionValidator{}
mw := newTunnelMiddleware(t, validator)
lookup := TunnelLookupFunc(func(_ netip.Addr) (PeerIdentity, bool) {
return PeerIdentity{}, false
})
w, r := newTunnelRequest("100.64.0.99:55555")
r = r.WithContext(WithTunnelLookup(r.Context(), lookup))
called := false
next := http.HandlerFunc(func(http.ResponseWriter, *http.Request) { called = true })
config, _ := mw.getDomainConfig("svc.example")
handled := mw.forwardWithTunnelPeer(w, r, "svc.example", config, next)
assert.False(t, handled, "unknown peer must fall through, not forward")
assert.False(t, called, "next handler must not run for unknown peer")
assert.Equal(t, int32(0), validator.tunnelCalls.Load(), "ValidateTunnelPeer must be skipped on local-lookup miss")
}
// TestForwardWithTunnelPeer_GroupsPropagateToCapturedData verifies the proxy
// surfaces the calling peer's group memberships from ValidateTunnelPeerResponse
// onto CapturedData so policy-aware middlewares can authorise without an
// extra management round-trip.
func TestForwardWithTunnelPeer_GroupsPropagateToCapturedData(t *testing.T) {
groups := []string{"engineering", "sre"}
validator := &stubSessionValidator{
respFn: func(_ *proto.ValidateTunnelPeerRequest) *proto.ValidateTunnelPeerResponse {
return &proto.ValidateTunnelPeerResponse{
Valid: true,
SessionToken: "tok",
UserId: "user-1",
PeerGroupIds: groups,
}
},
}
mw := newTunnelMiddleware(t, validator)
w, r := newTunnelRequest("100.64.0.10:55555")
cd := proxy.NewCapturedData("")
lookup := TunnelLookupFunc(func(_ netip.Addr) (PeerIdentity, bool) {
return PeerIdentity{}, true
})
r = r.WithContext(proxy.WithCapturedData(WithTunnelLookup(r.Context(), lookup), cd))
called := false
next := http.HandlerFunc(func(http.ResponseWriter, *http.Request) { called = true })
config, _ := mw.getDomainConfig("svc.example")
handled := mw.forwardWithTunnelPeer(w, r, "svc.example", config, next)
require.True(t, handled, "valid tunnel-peer response must forward")
require.True(t, called, "next handler must run")
assert.Equal(t, "user-1", cd.GetUserID(), "user id must propagate from tunnel-peer response")
assert.Equal(t, groups, cd.GetUserGroups(), "peer group IDs must propagate from tunnel-peer response")
}
// TestForwardWithTunnelPeer_LocalLookupKnownPeerStillRPCs verifies that a
// known tunnel IP still triggers ValidateTunnelPeer for the user-identity
// tail (UserID + group access). Phase 3 only short-circuits the deny path.
func TestForwardWithTunnelPeer_LocalLookupKnownPeerStillRPCs(t *testing.T) {
validator := &stubSessionValidator{
respFn: func(_ *proto.ValidateTunnelPeerRequest) *proto.ValidateTunnelPeerResponse {
return &proto.ValidateTunnelPeerResponse{Valid: true, SessionToken: "tok", UserId: "user-1"}
},
}
mw := newTunnelMiddleware(t, validator)
knownIP := netip.MustParseAddr("100.64.0.10")
lookup := TunnelLookupFunc(func(ip netip.Addr) (PeerIdentity, bool) {
if ip == knownIP {
return PeerIdentity{PubKey: "pk", TunnelIP: ip, FQDN: "peer.netbird.cloud"}, true
}
return PeerIdentity{}, false
})
w, r := newTunnelRequest(knownIP.String() + ":55555")
r = r.WithContext(WithTunnelLookup(r.Context(), lookup))
called := false
next := http.HandlerFunc(func(http.ResponseWriter, *http.Request) { called = true })
config, _ := mw.getDomainConfig("svc.example")
handled := mw.forwardWithTunnelPeer(w, r, "svc.example", config, next)
assert.True(t, handled, "known peer with valid RPC response must forward")
assert.True(t, called, "next handler must run on success")
assert.Equal(t, int32(1), validator.tunnelCalls.Load(), "RPC must run for the user-identity tail when local lookup confirms the peer")
}
// TestForwardWithTunnelPeer_NoLookupRefusesFastPath guards the
// anti-spoof gate: requests that didn't arrive on the per-account
// inbound listener (no TunnelLookup attached) must never reach
// management's ValidateTunnelPeer, even when the source IP looks like
// a tunnel address. A colliding RFC1918 / CGNAT source on the public
// listener would otherwise impersonate a mesh peer.
func TestForwardWithTunnelPeer_NoLookupRefusesFastPath(t *testing.T) {
validator := &stubSessionValidator{
respFn: func(_ *proto.ValidateTunnelPeerRequest) *proto.ValidateTunnelPeerResponse {
return &proto.ValidateTunnelPeerResponse{Valid: true, SessionToken: "tok", UserId: "user-1"}
},
}
mw := newTunnelMiddleware(t, validator)
w, r := newTunnelRequest("100.64.0.10:55555")
called := false
next := http.HandlerFunc(func(http.ResponseWriter, *http.Request) { called = true })
config, _ := mw.getDomainConfig("svc.example")
handled := mw.forwardWithTunnelPeer(w, r, "svc.example", config, next)
assert.False(t, handled, "fast-path must refuse without the inbound marker")
assert.False(t, called, "next handler must not run")
assert.Equal(t, int32(0), validator.tunnelCalls.Load(), "ValidateTunnelPeer must not be invoked without the inbound marker")
}
// TestForwardWithTunnelPeer_RPCErrorFallsThrough validates that an RPC
// failure still falls through to the next scheme (no false positive).
func TestForwardWithTunnelPeer_RPCErrorFallsThrough(t *testing.T) {
validator := &stubSessionValidator{respErr: errors.New("management down")}
mw := newTunnelMiddleware(t, validator)
knownIP := netip.MustParseAddr("100.64.0.10")
lookup := TunnelLookupFunc(func(ip netip.Addr) (PeerIdentity, bool) {
return PeerIdentity{TunnelIP: ip}, true
})
w, r := newTunnelRequest(knownIP.String() + ":55555")
r = r.WithContext(WithTunnelLookup(r.Context(), lookup))
config, _ := mw.getDomainConfig("svc.example")
handled := mw.forwardWithTunnelPeer(w, r, "svc.example", config, http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}))
assert.False(t, handled, "RPC error must let the caller try other schemes")
assert.Equal(t, int32(1), validator.tunnelCalls.Load(), "RPC was attempted exactly once")
}
// TestForwardWithTunnelPeer_CacheReusesPositiveResponse confirms the
// (account, IP, domain) cache prevents repeated RPCs for the same peer.
func TestForwardWithTunnelPeer_CacheReusesPositiveResponse(t *testing.T) {
validator := &stubSessionValidator{
respFn: func(_ *proto.ValidateTunnelPeerRequest) *proto.ValidateTunnelPeerResponse {
return &proto.ValidateTunnelPeerResponse{Valid: true, SessionToken: "tok", UserId: "user-1"}
},
}
mw := newTunnelMiddleware(t, validator)
lookup := TunnelLookupFunc(func(_ netip.Addr) (PeerIdentity, bool) {
return PeerIdentity{}, true
})
for i := 0; i < 4; i++ {
w, r := newTunnelRequest("100.64.0.10:55555")
r = r.WithContext(WithTunnelLookup(r.Context(), lookup))
next := http.HandlerFunc(func(http.ResponseWriter, *http.Request) {})
config, _ := mw.getDomainConfig("svc.example")
handled := mw.forwardWithTunnelPeer(w, r, "svc.example", config, next)
require.True(t, handled, "iteration %d should forward", i)
}
assert.Equal(t, int32(1), validator.tunnelCalls.Load(), "subsequent forwards must hit the cache, not management")
}
// TestForwardWithTunnelPeer_RoutesAccountIDIntoCacheKey ensures cache keys
// honour account scoping — same tunnel IP on different accounts must not
// collide.
func TestForwardWithTunnelPeer_RoutesAccountIDIntoCacheKey(t *testing.T) {
validator := &stubSessionValidator{
respFn: func(req *proto.ValidateTunnelPeerRequest) *proto.ValidateTunnelPeerResponse {
return &proto.ValidateTunnelPeerResponse{Valid: true, SessionToken: "tok", UserId: "user"}
},
}
mw := NewMiddleware(log.New(), validator, nil)
require.NoError(t, mw.AddDomain("svc-a.example", nil, "", 0, "acct-a", "svc-a", nil, false))
require.NoError(t, mw.AddDomain("svc-b.example", nil, "", 0, "acct-b", "svc-b", nil, false))
// The fast-path requires the inbound-listener marker on the context.
// The peerstore lookup itself is account-agnostic at this level
// (one TunnelLookupFunc per account is attached by inbound.go); a
// trivial "always hit" lookup is enough to exercise the cache-key
// branch this test covers.
lookup := TunnelLookupFunc(func(_ netip.Addr) (PeerIdentity, bool) {
return PeerIdentity{}, true
})
for _, host := range []string{"svc-a.example", "svc-b.example"} {
w := httptest.NewRecorder()
r := httptest.NewRequest(http.MethodGet, "https://"+host+"/", nil)
r.Host = host
r.RemoteAddr = "100.64.0.10:55555"
r = r.WithContext(WithTunnelLookup(r.Context(), lookup))
config, _ := mw.getDomainConfig(host)
handled := mw.forwardWithTunnelPeer(w, r, host, config, http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}))
require.True(t, handled, "host %s should forward", host)
}
assert.Equal(t, int32(2), validator.tunnelCalls.Load(), "cache must not collide across accounts even when tunnel IPs match")
}
// TestForwardWithTunnelPeer_LocalLookupShortCircuitDoesNotPopulateCache
// guarantees that the deny-fast path leaves the cache untouched, so a
// subsequent request from the same IP after the peerstore catches up
// goes through the normal RPC flow.
func TestForwardWithTunnelPeer_LocalLookupShortCircuitDoesNotPopulateCache(t *testing.T) {
validator := &stubSessionValidator{
respFn: func(_ *proto.ValidateTunnelPeerRequest) *proto.ValidateTunnelPeerResponse {
return &proto.ValidateTunnelPeerResponse{Valid: true, SessionToken: "tok"}
},
}
mw := newTunnelMiddleware(t, validator)
knownIP := netip.MustParseAddr("100.64.0.10")
known := false
lookup := TunnelLookupFunc(func(ip netip.Addr) (PeerIdentity, bool) {
if known && ip == knownIP {
return PeerIdentity{TunnelIP: ip}, true
}
return PeerIdentity{}, false
})
doRequest := func() bool {
w, r := newTunnelRequest(knownIP.String() + ":55555")
r = r.WithContext(WithTunnelLookup(r.Context(), lookup))
config, _ := mw.getDomainConfig("svc.example")
return mw.forwardWithTunnelPeer(w, r, "svc.example", config, http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}))
}
require.False(t, doRequest(), "first request must short-circuit")
require.Equal(t, int32(0), validator.tunnelCalls.Load(), "short-circuit must not populate the cache")
known = true
require.True(t, doRequest(), "second request with peer in roster must forward via RPC")
assert.Equal(t, int32(1), validator.tunnelCalls.Load(), "RPC runs once after peerstore catches up")
}
func TestPrivateService_FailsClosedOnTunnelPeerFailure(t *testing.T) {
mw := NewMiddleware(log.New(), nil, nil)
require.NoError(t, mw.AddDomain("private.svc", nil, "", 0, "acct-1", "svc-1", nil, true))
called := false
handler := mw.Protect(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
called = true
w.WriteHeader(http.StatusOK)
}))
req := httptest.NewRequest(http.MethodGet, "https://private.svc/", nil)
req.Host = "private.svc"
req.RemoteAddr = "100.64.0.10:55555"
w := httptest.NewRecorder()
handler.ServeHTTP(w, req)
assert.Equal(t, http.StatusForbidden, w.Code)
assert.False(t, called)
}
func TestPrivateService_ForwardsOnTunnelPeerSuccess(t *testing.T) {
validator := &stubSessionValidator{
respFn: func(_ *proto.ValidateTunnelPeerRequest) *proto.ValidateTunnelPeerResponse {
return &proto.ValidateTunnelPeerResponse{
Valid: true,
SessionToken: "tok",
UserId: "user-1",
}
},
}
mw := NewMiddleware(log.New(), validator, nil)
require.NoError(t, mw.AddDomain("private.svc", nil, "", 0, "acct-1", "svc-1", nil, true))
called := false
handler := mw.Protect(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
called = true
w.WriteHeader(http.StatusOK)
}))
// Per-account inbound listener attaches WithTunnelLookup; without it
// forwardWithTunnelPeer refuses to take the fast-path. Mirror the
// real flow so this test exercises the post-gating success branch.
lookup := TunnelLookupFunc(func(_ netip.Addr) (PeerIdentity, bool) {
return PeerIdentity{}, true
})
req := httptest.NewRequest(http.MethodGet, "https://private.svc/", nil)
req.Host = "private.svc"
req.RemoteAddr = "100.64.0.10:55555"
req = req.WithContext(WithTunnelLookup(req.Context(), lookup))
w := httptest.NewRecorder()
handler.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
assert.True(t, called)
}
proxy/internal/certwatch/watcher.go
+286
View File
@@ -0,0 +1,286 @@
// Package certwatch watches TLS certificate files on disk and provides
// a hot-reloading GetCertificate callback for tls.Config.
package certwatch
import (
"context"
"crypto/tls"
"crypto/x509"
"fmt"
"path/filepath"
"sync"
"time"
"github.com/fsnotify/fsnotify"
log "github.com/sirupsen/logrus"
)
const (
defaultPollInterval = 30 * time.Second
debounceDelay = 500 * time.Millisecond
)
// Watcher monitors TLS certificate files on disk and caches the loaded
// certificate in memory. It detects changes via fsnotify (with a polling
// fallback for filesystems like NFS that lack inotify support) and
// reloads the certificate pair automatically.
type Watcher struct {
certPath string
keyPath string
mu sync.RWMutex
cert *tls.Certificate
leaf *x509.Certificate
pollInterval time.Duration
logger *log.Logger
}
// NewWatcher creates a Watcher that monitors the given cert and key files.
// It performs an initial load of the certificate and returns an error
// if the initial load fails.
func NewWatcher(certPath, keyPath string, logger *log.Logger) (*Watcher, error) {
if logger == nil {
logger = log.StandardLogger()
}
w := &Watcher{
certPath: certPath,
keyPath: keyPath,
pollInterval: defaultPollInterval,
logger: logger,
}
if err := w.reload(); err != nil {
return nil, fmt.Errorf("initial certificate load: %w", err)
}
return w, nil
}
// GetCertificate returns the current in-memory certificate.
// It is safe for concurrent use and compatible with tls.Config.GetCertificate.
func (w *Watcher) GetCertificate(_ *tls.ClientHelloInfo) (*tls.Certificate, error) {
w.mu.RLock()
defer w.mu.RUnlock()
return w.cert, nil
}
// Leaf returns the parsed leaf certificate, or nil if not yet loaded.
func (w *Watcher) Leaf() *x509.Certificate {
w.mu.RLock()
defer w.mu.RUnlock()
return w.leaf
}
// Watch starts watching for certificate file changes. It blocks until
// ctx is cancelled. It uses fsnotify for immediate detection and falls
// back to polling if fsnotify is unavailable (e.g. on NFS).
// Even with fsnotify active, a periodic poll runs as a safety net.
func (w *Watcher) Watch(ctx context.Context) {
// Watch the parent directory rather than individual files. Some volume
// mounts use an atomic symlink swap (..data -> timestamped dir), so
// watching the parent directory catches the link replacement.
certDir := filepath.Dir(w.certPath)
keyDir := filepath.Dir(w.keyPath)
watcher, err := fsnotify.NewWatcher()
if err != nil {
w.logger.Warnf("fsnotify unavailable, using polling only: %v", err)
w.pollLoop(ctx)
return
}
defer func() {
if err := watcher.Close(); err != nil {
w.logger.Debugf("close fsnotify watcher: %v", err)
}
}()
if err := watcher.Add(certDir); err != nil {
w.logger.Warnf("fsnotify watch on %s failed, using polling only: %v", certDir, err)
w.pollLoop(ctx)
return
}
if keyDir != certDir {
if err := watcher.Add(keyDir); err != nil {
w.logger.Warnf("fsnotify watch on %s failed: %v", keyDir, err)
}
}
w.logger.Infof("watching certificate files in %s", certDir)
w.fsnotifyLoop(ctx, watcher)
}
func (w *Watcher) fsnotifyLoop(ctx context.Context, watcher *fsnotify.Watcher) {
certBase := filepath.Base(w.certPath)
keyBase := filepath.Base(w.keyPath)
var debounce *time.Timer
defer func() {
if debounce != nil {
debounce.Stop()
}
}()
// Periodic poll as a safety net for missed fsnotify events.
pollTicker := time.NewTicker(w.pollInterval)
defer pollTicker.Stop()
for {
select {
case <-ctx.Done():
return
case event, ok := <-watcher.Events:
if !ok {
return
}
base := filepath.Base(event.Name)
if !isRelevantFile(base, certBase, keyBase) {
w.logger.Debugf("fsnotify: ignoring event %s on %s", event.Op, event.Name)
continue
}
if !event.Has(fsnotify.Create) && !event.Has(fsnotify.Write) && !event.Has(fsnotify.Rename) {
w.logger.Debugf("fsnotify: ignoring op %s on %s", event.Op, base)
continue
}
w.logger.Debugf("fsnotify: detected %s on %s, scheduling reload", event.Op, base)
// Debounce: cert-manager may write cert and key as separate
// operations. Wait briefly to load both at once.
if debounce != nil {
debounce.Stop()
}
debounce = time.AfterFunc(debounceDelay, func() {
if ctx.Err() != nil {
return
}
w.tryReload()
})
case err, ok := <-watcher.Errors:
if !ok {
return
}
w.logger.Warnf("fsnotify error: %v", err)
case <-pollTicker.C:
w.tryReload()
}
}
}
func (w *Watcher) pollLoop(ctx context.Context) {
ticker := time.NewTicker(w.pollInterval)
defer ticker.Stop()
for {
select {
case <-ctx.Done():
return
case <-ticker.C:
w.tryReload()
}
}
}
// reload loads the certificate from disk and updates the in-memory cache.
func (w *Watcher) reload() error {
cert, err := tls.LoadX509KeyPair(w.certPath, w.keyPath)
if err != nil {
return err
}
// Parse the leaf for comparison on subsequent reloads.
if cert.Leaf == nil && len(cert.Certificate) > 0 {
leaf, err := x509.ParseCertificate(cert.Certificate[0])
if err != nil {
return fmt.Errorf("parse leaf certificate: %w", err)
}
cert.Leaf = leaf
}
w.mu.Lock()
w.cert = &cert
w.leaf = cert.Leaf
w.mu.Unlock()
w.logCertDetails("loaded certificate", cert.Leaf)
return nil
}
// tryReload attempts to reload the certificate. It skips the update
// if the certificate on disk is identical to the one in memory (same
// serial number and issuer) to avoid redundant log noise.
func (w *Watcher) tryReload() {
cert, err := tls.LoadX509KeyPair(w.certPath, w.keyPath)
if err != nil {
w.logger.Warnf("reload certificate: %v", err)
return
}
if cert.Leaf == nil && len(cert.Certificate) > 0 {
leaf, err := x509.ParseCertificate(cert.Certificate[0])
if err != nil {
w.logger.Warnf("parse reloaded leaf certificate: %v", err)
return
}
cert.Leaf = leaf
}
w.mu.Lock()
if w.leaf != nil && cert.Leaf != nil &&
w.leaf.SerialNumber.Cmp(cert.Leaf.SerialNumber) == 0 &&
w.leaf.Issuer.CommonName == cert.Leaf.Issuer.CommonName {
w.mu.Unlock()
return
}
prev := w.leaf
w.cert = &cert
w.leaf = cert.Leaf
w.mu.Unlock()
w.logCertChange(prev, cert.Leaf)
}
func (w *Watcher) logCertDetails(msg string, leaf *x509.Certificate) {
if leaf == nil {
w.logger.Info(msg)
return
}
w.logger.Infof("%s: subject=%q serial=%s SANs=%v notAfter=%s",
msg,
leaf.Subject.CommonName,
leaf.SerialNumber.Text(16),
leaf.DNSNames,
leaf.NotAfter.UTC().Format(time.RFC3339),
)
}
func (w *Watcher) logCertChange(prev, next *x509.Certificate) {
if prev == nil || next == nil {
w.logCertDetails("certificate reloaded from disk", next)
return
}
w.logger.Infof("certificate reloaded from disk: subject=%q -> %q serial=%s -> %s notAfter=%s -> %s",
prev.Subject.CommonName, next.Subject.CommonName,
prev.SerialNumber.Text(16), next.SerialNumber.Text(16),
prev.NotAfter.UTC().Format(time.RFC3339), next.NotAfter.UTC().Format(time.RFC3339),
)
}
// isRelevantFile returns true if the changed file name is one we care about.
// This includes the cert/key files themselves and the ..data symlink used
// by atomic volume mounts.
func isRelevantFile(changed, certBase, keyBase string) bool {
return changed == certBase || changed == keyBase || changed == "..data"
}
proxy/internal/certwatch/watcher_test.go
+292
View File
@@ -0,0 +1,292 @@
package certwatch
import (
"context"
"crypto/ecdsa"
"crypto/elliptic"
"crypto/rand"
"crypto/tls"
"crypto/x509"
"crypto/x509/pkix"
"encoding/pem"
"math/big"
"os"
"path/filepath"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func generateSelfSignedCert(t *testing.T, serial int64) (certPEM, keyPEM []byte) {
t.Helper()
key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
require.NoError(t, err)
template := &x509.Certificate{
SerialNumber: big.NewInt(serial),
Subject: pkix.Name{CommonName: "test"},
NotBefore: time.Now().Add(-time.Hour),
NotAfter: time.Now().Add(24 * time.Hour),
}
certDER, err := x509.CreateCertificate(rand.Reader, template, template, &key.PublicKey, key)
require.NoError(t, err)
certPEM = pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: certDER})
keyDER, err := x509.MarshalECPrivateKey(key)
require.NoError(t, err)
keyPEM = pem.EncodeToMemory(&pem.Block{Type: "EC PRIVATE KEY", Bytes: keyDER})
return certPEM, keyPEM
}
func writeCert(t *testing.T, dir string, certPEM, keyPEM []byte) {
t.Helper()
require.NoError(t, os.WriteFile(filepath.Join(dir, "tls.crt"), certPEM, 0o600))
require.NoError(t, os.WriteFile(filepath.Join(dir, "tls.key"), keyPEM, 0o600))
}
func TestNewWatcher(t *testing.T) {
dir := t.TempDir()
certPEM, keyPEM := generateSelfSignedCert(t, 1)
writeCert(t, dir, certPEM, keyPEM)
w, err := NewWatcher(
filepath.Join(dir, "tls.crt"),
filepath.Join(dir, "tls.key"),
nil,
)
require.NoError(t, err)
cert, err := w.GetCertificate(nil)
require.NoError(t, err)
require.NotNil(t, cert)
assert.Equal(t, int64(1), cert.Leaf.SerialNumber.Int64())
}
func TestNewWatcherMissingFiles(t *testing.T) {
dir := t.TempDir()
_, err := NewWatcher(
filepath.Join(dir, "tls.crt"),
filepath.Join(dir, "tls.key"),
nil,
)
assert.Error(t, err)
}
func TestReload(t *testing.T) {
dir := t.TempDir()
certPEM1, keyPEM1 := generateSelfSignedCert(t, 100)
writeCert(t, dir, certPEM1, keyPEM1)
w, err := NewWatcher(
filepath.Join(dir, "tls.crt"),
filepath.Join(dir, "tls.key"),
nil,
)
require.NoError(t, err)
cert1, err := w.GetCertificate(nil)
require.NoError(t, err)
assert.Equal(t, int64(100), cert1.Leaf.SerialNumber.Int64())
// Write a new cert with a different serial.
certPEM2, keyPEM2 := generateSelfSignedCert(t, 200)
writeCert(t, dir, certPEM2, keyPEM2)
// Manually trigger reload.
w.tryReload()
cert2, err := w.GetCertificate(nil)
require.NoError(t, err)
assert.Equal(t, int64(200), cert2.Leaf.SerialNumber.Int64())
}
func TestTryReloadSkipsUnchanged(t *testing.T) {
dir := t.TempDir()
certPEM, keyPEM := generateSelfSignedCert(t, 42)
writeCert(t, dir, certPEM, keyPEM)
w, err := NewWatcher(
filepath.Join(dir, "tls.crt"),
filepath.Join(dir, "tls.key"),
nil,
)
require.NoError(t, err)
cert1, err := w.GetCertificate(nil)
require.NoError(t, err)
// Reload with same cert - pointer should remain the same.
w.tryReload()
cert2, err := w.GetCertificate(nil)
require.NoError(t, err)
assert.Same(t, cert1, cert2, "cert pointer should not change when content is the same")
}
func TestWatchDetectsChanges(t *testing.T) {
dir := t.TempDir()
certPEM1, keyPEM1 := generateSelfSignedCert(t, 1)
writeCert(t, dir, certPEM1, keyPEM1)
w, err := NewWatcher(
filepath.Join(dir, "tls.crt"),
filepath.Join(dir, "tls.key"),
nil,
)
require.NoError(t, err)
// Use a short poll interval for the test.
w.pollInterval = 100 * time.Millisecond
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
go w.Watch(ctx)
// Write new cert.
certPEM2, keyPEM2 := generateSelfSignedCert(t, 999)
writeCert(t, dir, certPEM2, keyPEM2)
// Wait for the watcher to pick it up.
require.Eventually(t, func() bool {
cert, err := w.GetCertificate(nil)
if err != nil {
return false
}
return cert.Leaf.SerialNumber.Int64() == 999
}, 5*time.Second, 50*time.Millisecond, "watcher should detect cert change")
}
func TestIsRelevantFile(t *testing.T) {
assert.True(t, isRelevantFile("tls.crt", "tls.crt", "tls.key"))
assert.True(t, isRelevantFile("tls.key", "tls.crt", "tls.key"))
assert.True(t, isRelevantFile("..data", "tls.crt", "tls.key"))
assert.False(t, isRelevantFile("other.txt", "tls.crt", "tls.key"))
}
// TestWatchSymlinkRotation simulates Kubernetes secret volume updates where
// the data directory is atomically swapped via a ..data symlink.
func TestWatchSymlinkRotation(t *testing.T) {
base := t.TempDir()
// Create initial target directory with certs.
dir1 := filepath.Join(base, "dir1")
require.NoError(t, os.Mkdir(dir1, 0o755))
certPEM1, keyPEM1 := generateSelfSignedCert(t, 1)
require.NoError(t, os.WriteFile(filepath.Join(dir1, "tls.crt"), certPEM1, 0o600))
require.NoError(t, os.WriteFile(filepath.Join(dir1, "tls.key"), keyPEM1, 0o600))
// Create ..data symlink pointing to dir1.
dataLink := filepath.Join(base, "..data")
require.NoError(t, os.Symlink(dir1, dataLink))
// Create tls.crt and tls.key as symlinks to ..data/{file}.
certLink := filepath.Join(base, "tls.crt")
keyLink := filepath.Join(base, "tls.key")
require.NoError(t, os.Symlink(filepath.Join(dataLink, "tls.crt"), certLink))
require.NoError(t, os.Symlink(filepath.Join(dataLink, "tls.key"), keyLink))
w, err := NewWatcher(certLink, keyLink, nil)
require.NoError(t, err)
cert, err := w.GetCertificate(nil)
require.NoError(t, err)
assert.Equal(t, int64(1), cert.Leaf.SerialNumber.Int64())
w.pollInterval = 100 * time.Millisecond
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
go w.Watch(ctx)
// Simulate k8s atomic rotation: create dir2, swap ..data symlink.
dir2 := filepath.Join(base, "dir2")
require.NoError(t, os.Mkdir(dir2, 0o755))
certPEM2, keyPEM2 := generateSelfSignedCert(t, 777)
require.NoError(t, os.WriteFile(filepath.Join(dir2, "tls.crt"), certPEM2, 0o600))
require.NoError(t, os.WriteFile(filepath.Join(dir2, "tls.key"), keyPEM2, 0o600))
// Atomic swap: create temp link, then rename over ..data.
tmpLink := filepath.Join(base, "..data_tmp")
require.NoError(t, os.Symlink(dir2, tmpLink))
require.NoError(t, os.Rename(tmpLink, dataLink))
require.Eventually(t, func() bool {
cert, err := w.GetCertificate(nil)
if err != nil {
return false
}
return cert.Leaf.SerialNumber.Int64() == 777
}, 5*time.Second, 50*time.Millisecond, "watcher should detect symlink rotation")
}
// TestPollLoopDetectsChanges verifies the poll-only fallback path works.
func TestPollLoopDetectsChanges(t *testing.T) {
dir := t.TempDir()
certPEM1, keyPEM1 := generateSelfSignedCert(t, 1)
writeCert(t, dir, certPEM1, keyPEM1)
w, err := NewWatcher(
filepath.Join(dir, "tls.crt"),
filepath.Join(dir, "tls.key"),
nil,
)
require.NoError(t, err)
w.pollInterval = 100 * time.Millisecond
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
// Directly use pollLoop to test the fallback path.
go w.pollLoop(ctx)
certPEM2, keyPEM2 := generateSelfSignedCert(t, 555)
writeCert(t, dir, certPEM2, keyPEM2)
require.Eventually(t, func() bool {
cert, err := w.GetCertificate(nil)
if err != nil {
return false
}
return cert.Leaf.SerialNumber.Int64() == 555
}, 5*time.Second, 50*time.Millisecond, "poll loop should detect cert change")
}
func TestGetCertificateConcurrency(t *testing.T) {
dir := t.TempDir()
certPEM, keyPEM := generateSelfSignedCert(t, 1)
writeCert(t, dir, certPEM, keyPEM)
w, err := NewWatcher(
filepath.Join(dir, "tls.crt"),
filepath.Join(dir, "tls.key"),
nil,
)
require.NoError(t, err)
// Hammer GetCertificate concurrently while reloading.
done := make(chan struct{})
go func() {
for i := 0; i < 100; i++ {
w.tryReload()
}
close(done)
}()
for i := 0; i < 1000; i++ {
cert, err := w.GetCertificate(&tls.ClientHelloInfo{})
assert.NoError(t, err)
assert.NotNil(t, cert)
}
<-done
}
proxy/internal/conntrack/conn.go
+51
View File
@@ -0,0 +1,51 @@
package conntrack
import (
"bufio"
"net"
"net/http"
)
// trackedConn wraps a net.Conn and removes itself from the tracker on Close.
type trackedConn struct {
net.Conn
tracker *HijackTracker
host string
}
func (c *trackedConn) Close() error {
c.tracker.remove(c)
return c.Conn.Close()
}
// trackingWriter wraps an http.ResponseWriter and intercepts Hijack calls
// to replace the raw connection with a trackedConn that auto-deregisters.
type trackingWriter struct {
http.ResponseWriter
tracker *HijackTracker
host string
}
func (w *trackingWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
hijacker, ok := w.ResponseWriter.(http.Hijacker)
if !ok {
return nil, nil, http.ErrNotSupported
}
conn, buf, err := hijacker.Hijack()
if err != nil {
return nil, nil, err
}
tc := &trackedConn{Conn: conn, tracker: w.tracker, host: w.host}
w.tracker.add(tc)
return tc, buf, nil
}
func (w *trackingWriter) Flush() {
if flusher, ok := w.ResponseWriter.(http.Flusher); ok {
flusher.Flush()
}
}
func (w *trackingWriter) Unwrap() http.ResponseWriter {
return w.ResponseWriter
}
proxy/internal/conntrack/hijacked.go
+96
View File
@@ -0,0 +1,96 @@
package conntrack
import (
"net/http"
"sync"
)
// HijackTracker tracks connections that have been hijacked (e.g. WebSocket
// upgrades). http.Server.Shutdown does not close hijacked connections, so
// they must be tracked and closed explicitly during graceful shutdown.
//
// Connections are indexed by the request Host so they can be closed
// per-domain when a service mapping is removed.
//
// Use Middleware as the outermost HTTP middleware to ensure hijacked
// connections are tracked and automatically deregistered when closed.
type HijackTracker struct {
mu sync.Mutex
conns map[*trackedConn]struct{}
}
// Middleware returns an HTTP middleware that wraps the ResponseWriter so that
// hijacked connections are tracked and automatically deregistered from the
// tracker when closed. This should be the outermost middleware in the chain.
func (t *HijackTracker) Middleware(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
next.ServeHTTP(&trackingWriter{
ResponseWriter: w,
tracker: t,
host: hostOnly(r.Host),
}, r)
})
}
// CloseAll closes all tracked hijacked connections and returns the count.
func (t *HijackTracker) CloseAll() int {
t.mu.Lock()
conns := t.conns
t.conns = nil
t.mu.Unlock()
for tc := range conns {
_ = tc.Conn.Close()
}
return len(conns)
}
// CloseByHost closes all tracked hijacked connections for the given host
// and returns the number of connections closed.
func (t *HijackTracker) CloseByHost(host string) int {
host = hostOnly(host)
t.mu.Lock()
var toClose []*trackedConn
for tc := range t.conns {
if tc.host == host {
toClose = append(toClose, tc)
}
}
for _, tc := range toClose {
delete(t.conns, tc)
}
t.mu.Unlock()
for _, tc := range toClose {
_ = tc.Conn.Close()
}
return len(toClose)
}
func (t *HijackTracker) add(tc *trackedConn) {
t.mu.Lock()
if t.conns == nil {
t.conns = make(map[*trackedConn]struct{})
}
t.conns[tc] = struct{}{}
t.mu.Unlock()
}
func (t *HijackTracker) remove(tc *trackedConn) {
t.mu.Lock()
delete(t.conns, tc)
t.mu.Unlock()
}
// hostOnly strips the port from a host:port string.
func hostOnly(hostport string) string {
for i := len(hostport) - 1; i >= 0; i-- {
if hostport[i] == ':' {
return hostport[:i]
}
if hostport[i] < '0' || hostport[i] > '9' {
return hostport
}
}
return hostport
}
proxy/internal/conntrack/hijacked_test.go
+142
View File
@@ -0,0 +1,142 @@
package conntrack
import (
"bufio"
"net"
"net/http"
"net/http/httptest"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
// fakeHijackWriter implements http.ResponseWriter and http.Hijacker for testing.
type fakeHijackWriter struct {
http.ResponseWriter
conn net.Conn
}
func (f *fakeHijackWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
rw := bufio.NewReadWriter(bufio.NewReader(f.conn), bufio.NewWriter(f.conn))
return f.conn, rw, nil
}
func TestCloseByHost(t *testing.T) {
var tracker HijackTracker
// Simulate hijacking two connections for different hosts.
connA1, connA2 := net.Pipe()
defer connA2.Close()
connB1, connB2 := net.Pipe()
defer connB2.Close()
twA := &trackingWriter{
ResponseWriter: httptest.NewRecorder(),
tracker: &tracker,
host: "a.example.com",
}
twB := &trackingWriter{
ResponseWriter: httptest.NewRecorder(),
tracker: &tracker,
host: "b.example.com",
}
// Use fakeHijackWriter to provide the Hijack method.
twA.ResponseWriter = &fakeHijackWriter{ResponseWriter: twA.ResponseWriter, conn: connA1}
twB.ResponseWriter = &fakeHijackWriter{ResponseWriter: twB.ResponseWriter, conn: connB1}
_, _, err := twA.Hijack()
require.NoError(t, err)
_, _, err = twB.Hijack()
require.NoError(t, err)
tracker.mu.Lock()
assert.Equal(t, 2, len(tracker.conns), "should track 2 connections")
tracker.mu.Unlock()
// Close only host A.
n := tracker.CloseByHost("a.example.com")
assert.Equal(t, 1, n, "should close 1 connection for host A")
tracker.mu.Lock()
assert.Equal(t, 1, len(tracker.conns), "should have 1 remaining connection")
tracker.mu.Unlock()
// Verify host A's conn is actually closed.
buf := make([]byte, 1)
_, err = connA2.Read(buf)
assert.Error(t, err, "host A pipe should be closed")
// Host B should still be alive.
go func() { _, _ = connB1.Write([]byte("x")) }()
// Close all remaining.
n = tracker.CloseAll()
assert.Equal(t, 1, n, "should close remaining 1 connection")
tracker.mu.Lock()
assert.Equal(t, 0, len(tracker.conns), "should have 0 connections after CloseAll")
tracker.mu.Unlock()
}
func TestCloseAll(t *testing.T) {
var tracker HijackTracker
for range 5 {
c1, c2 := net.Pipe()
defer c2.Close()
tc := &trackedConn{Conn: c1, tracker: &tracker, host: "test.com"}
tracker.add(tc)
}
tracker.mu.Lock()
assert.Equal(t, 5, len(tracker.conns))
tracker.mu.Unlock()
n := tracker.CloseAll()
assert.Equal(t, 5, n)
// Double CloseAll is safe.
n = tracker.CloseAll()
assert.Equal(t, 0, n)
}
func TestTrackedConn_AutoDeregister(t *testing.T) {
var tracker HijackTracker
c1, c2 := net.Pipe()
defer c2.Close()
tc := &trackedConn{Conn: c1, tracker: &tracker, host: "auto.com"}
tracker.add(tc)
tracker.mu.Lock()
assert.Equal(t, 1, len(tracker.conns))
tracker.mu.Unlock()
// Close the tracked conn: should auto-deregister.
require.NoError(t, tc.Close())
tracker.mu.Lock()
assert.Equal(t, 0, len(tracker.conns), "should auto-deregister on close")
tracker.mu.Unlock()
}
func TestHostOnly(t *testing.T) {
tests := []struct {
input string
want string
}{
{"example.com:443", "example.com"},
{"example.com", "example.com"},
{"127.0.0.1:8080", "127.0.0.1"},
{"[::1]:443", "[::1]"},
{"", ""},
}
for _, tt := range tests {
t.Run(tt.input, func(t *testing.T) {
assert.Equal(t, tt.want, hostOnly(tt.input))
})
}
}
proxy/internal/crowdsec/bouncer.go
+251
View File
@@ -0,0 +1,251 @@
// Package crowdsec provides a CrowdSec stream bouncer that maintains a local
// decision cache for IP reputation checks.
package crowdsec
import (
"context"
"errors"
"net/netip"
"strings"
"sync"
"sync/atomic"
"time"
"github.com/crowdsecurity/crowdsec/pkg/models"
csbouncer "github.com/crowdsecurity/go-cs-bouncer"
log "github.com/sirupsen/logrus"
"github.com/netbirdio/netbird/proxy/internal/restrict"
)
// Bouncer wraps a CrowdSec StreamBouncer, maintaining a local cache of
// active decisions for fast IP lookups. It implements restrict.CrowdSecChecker.
type Bouncer struct {
mu sync.RWMutex
ips map[netip.Addr]*restrict.CrowdSecDecision
prefixes map[netip.Prefix]*restrict.CrowdSecDecision
ready atomic.Bool
apiURL string
apiKey string
tickerInterval time.Duration
logger *log.Entry
// lifeMu protects cancel and done from concurrent Start/Stop calls.
lifeMu sync.Mutex
cancel context.CancelFunc
done chan struct{}
}
// compile-time check
var _ restrict.CrowdSecChecker = (*Bouncer)(nil)
// NewBouncer creates a bouncer but does not start the stream.
func NewBouncer(apiURL, apiKey string, logger *log.Entry) *Bouncer {
return &Bouncer{
apiURL: apiURL,
apiKey: apiKey,
logger: logger,
ips: make(map[netip.Addr]*restrict.CrowdSecDecision),
prefixes: make(map[netip.Prefix]*restrict.CrowdSecDecision),
}
}
// Start launches the background goroutine that streams decisions from the
// CrowdSec LAPI. The stream runs until Stop is called or ctx is cancelled.
func (b *Bouncer) Start(ctx context.Context) error {
interval := b.tickerInterval
if interval == 0 {
interval = 10 * time.Second
}
stream := &csbouncer.StreamBouncer{
APIKey: b.apiKey,
APIUrl: b.apiURL,
TickerInterval: interval.String(),
UserAgent: "netbird-proxy/1.0",
Scopes: []string{"ip", "range"},
RetryInitialConnect: true,
}
b.logger.Infof("connecting to CrowdSec LAPI at %s", b.apiURL)
if err := stream.Init(); err != nil {
return err
}
// Reset state from any previous run.
b.mu.Lock()
b.ips = make(map[netip.Addr]*restrict.CrowdSecDecision)
b.prefixes = make(map[netip.Prefix]*restrict.CrowdSecDecision)
b.mu.Unlock()
b.ready.Store(false)
ctx, cancel := context.WithCancel(ctx)
done := make(chan struct{})
b.lifeMu.Lock()
if b.cancel != nil {
b.lifeMu.Unlock()
cancel()
return errors.New("bouncer already started")
}
b.cancel = cancel
b.done = done
b.lifeMu.Unlock()
var wg sync.WaitGroup
wg.Add(2)
go func() {
defer wg.Done()
if err := stream.Run(ctx); err != nil && ctx.Err() == nil {
b.logger.Errorf("CrowdSec stream ended: %v", err)
}
}()
go func() {
defer wg.Done()
b.consumeStream(ctx, stream)
}()
go func() {
wg.Wait()
close(done)
}()
return nil
}
// Stop cancels the stream and waits for all goroutines to finish.
func (b *Bouncer) Stop() {
b.lifeMu.Lock()
cancel := b.cancel
done := b.done
b.cancel = nil
b.lifeMu.Unlock()
if cancel != nil {
cancel()
<-done
}
}
// Ready returns true after the first batch of decisions has been processed.
func (b *Bouncer) Ready() bool {
return b.ready.Load()
}
// CheckIP looks up addr in the local decision cache. Returns nil if no
// active decision exists for the address.
//
// Prefix lookups are O(1): instead of scanning all stored prefixes, we
// probe the map for every possible containing prefix of the address
// (at most 33 for IPv4, 129 for IPv6).
func (b *Bouncer) CheckIP(addr netip.Addr) *restrict.CrowdSecDecision {
addr = addr.Unmap()
b.mu.RLock()
defer b.mu.RUnlock()
if d, ok := b.ips[addr]; ok {
return d
}
maxBits := 32
if addr.Is6() {
maxBits = 128
}
// Walk from most-specific to least-specific prefix so the narrowest
// matching decision wins when ranges overlap.
for bits := maxBits; bits >= 0; bits-- {
prefix := netip.PrefixFrom(addr, bits).Masked()
if d, ok := b.prefixes[prefix]; ok {
return d
}
}
return nil
}
func (b *Bouncer) consumeStream(ctx context.Context, stream *csbouncer.StreamBouncer) {
first := true
for {
select {
case <-ctx.Done():
return
case resp, ok := <-stream.Stream:
if !ok {
return
}
b.mu.Lock()
b.applyDeleted(resp.Deleted)
b.applyNew(resp.New)
b.mu.Unlock()
if first {
b.ready.Store(true)
b.logger.Info("CrowdSec bouncer synced initial decisions")
first = false
}
}
}
}
func (b *Bouncer) applyDeleted(decisions []*models.Decision) {
for _, d := range decisions {
if d.Value == nil || d.Scope == nil {
continue
}
value := *d.Value
if strings.ToLower(*d.Scope) == "range" || strings.Contains(value, "/") {
prefix, err := netip.ParsePrefix(value)
if err != nil {
b.logger.Debugf("skip unparsable CrowdSec range deletion %q: %v", value, err)
continue
}
prefix = normalizePrefix(prefix)
delete(b.prefixes, prefix)
} else {
addr, err := netip.ParseAddr(value)
if err != nil {
b.logger.Debugf("skip unparsable CrowdSec IP deletion %q: %v", value, err)
continue
}
delete(b.ips, addr.Unmap())
}
}
}
func (b *Bouncer) applyNew(decisions []*models.Decision) {
for _, d := range decisions {
if d.Value == nil || d.Type == nil || d.Scope == nil {
continue
}
dec := &restrict.CrowdSecDecision{Type: restrict.DecisionType(*d.Type)}
value := *d.Value
if strings.ToLower(*d.Scope) == "range" || strings.Contains(value, "/") {
prefix, err := netip.ParsePrefix(value)
if err != nil {
b.logger.Debugf("skip unparsable CrowdSec range %q: %v", value, err)
continue
}
prefix = normalizePrefix(prefix)
b.prefixes[prefix] = dec
} else {
addr, err := netip.ParseAddr(value)
if err != nil {
b.logger.Debugf("skip unparsable CrowdSec IP %q: %v", value, err)
continue
}
b.ips[addr.Unmap()] = dec
}
}
}
// normalizePrefix unmaps v4-mapped-v6 addresses and zeros host bits so
// the prefix is a valid map key that matches CheckIP's probe logic.
func normalizePrefix(p netip.Prefix) netip.Prefix {
return netip.PrefixFrom(p.Addr().Unmap(), p.Bits()).Masked()
}
proxy/internal/crowdsec/bouncer_test.go
+337
View File
@@ -0,0 +1,337 @@
package crowdsec
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"net/netip"
"sync"
"testing"
"time"
"github.com/crowdsecurity/crowdsec/pkg/models"
log "github.com/sirupsen/logrus"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/netbirdio/netbird/proxy/internal/restrict"
)
func TestBouncer_CheckIP_Empty(t *testing.T) {
b := newTestBouncer()
b.ready.Store(true)
assert.Nil(t, b.CheckIP(netip.MustParseAddr("1.2.3.4")))
}
func TestBouncer_CheckIP_ExactMatch(t *testing.T) {
b := newTestBouncer()
b.ready.Store(true)
b.ips[netip.MustParseAddr("10.0.0.1")] = &restrict.CrowdSecDecision{Type: restrict.DecisionBan}
d := b.CheckIP(netip.MustParseAddr("10.0.0.1"))
require.NotNil(t, d)
assert.Equal(t, restrict.DecisionBan, d.Type)
assert.Nil(t, b.CheckIP(netip.MustParseAddr("10.0.0.2")))
}
func TestBouncer_CheckIP_PrefixMatch(t *testing.T) {
b := newTestBouncer()
b.ready.Store(true)
b.prefixes[netip.MustParsePrefix("192.168.1.0/24")] = &restrict.CrowdSecDecision{Type: restrict.DecisionBan}
d := b.CheckIP(netip.MustParseAddr("192.168.1.100"))
require.NotNil(t, d)
assert.Equal(t, restrict.DecisionBan, d.Type)
assert.Nil(t, b.CheckIP(netip.MustParseAddr("192.168.2.1")))
}
func TestBouncer_CheckIP_UnmapsV4InV6(t *testing.T) {
b := newTestBouncer()
b.ready.Store(true)
b.ips[netip.MustParseAddr("10.0.0.1")] = &restrict.CrowdSecDecision{Type: restrict.DecisionBan}
d := b.CheckIP(netip.MustParseAddr("::ffff:10.0.0.1"))
require.NotNil(t, d)
assert.Equal(t, restrict.DecisionBan, d.Type)
}
func TestBouncer_Ready(t *testing.T) {
b := newTestBouncer()
assert.False(t, b.Ready())
b.ready.Store(true)
assert.True(t, b.Ready())
}
func TestBouncer_CheckIP_ExactBeforePrefix(t *testing.T) {
b := newTestBouncer()
b.ready.Store(true)
b.ips[netip.MustParseAddr("10.0.0.1")] = &restrict.CrowdSecDecision{Type: restrict.DecisionCaptcha}
b.prefixes[netip.MustParsePrefix("10.0.0.0/8")] = &restrict.CrowdSecDecision{Type: restrict.DecisionBan}
d := b.CheckIP(netip.MustParseAddr("10.0.0.1"))
require.NotNil(t, d)
assert.Equal(t, restrict.DecisionCaptcha, d.Type)
d2 := b.CheckIP(netip.MustParseAddr("10.0.0.2"))
require.NotNil(t, d2)
assert.Equal(t, restrict.DecisionBan, d2.Type)
}
func TestBouncer_ApplyNew_IP(t *testing.T) {
b := newTestBouncer()
b.applyNew(makeDecisions(
decision{scope: "ip", value: "1.2.3.4", dtype: "ban", scenario: "test/brute"},
decision{scope: "ip", value: "5.6.7.8", dtype: "captcha", scenario: "test/crawl"},
))
require.Len(t, b.ips, 2)
assert.Equal(t, restrict.DecisionBan, b.ips[netip.MustParseAddr("1.2.3.4")].Type)
assert.Equal(t, restrict.DecisionCaptcha, b.ips[netip.MustParseAddr("5.6.7.8")].Type)
}
func TestBouncer_ApplyNew_Range(t *testing.T) {
b := newTestBouncer()
b.applyNew(makeDecisions(
decision{scope: "range", value: "10.0.0.0/8", dtype: "ban"},
))
require.Len(t, b.prefixes, 1)
assert.NotNil(t, b.prefixes[netip.MustParsePrefix("10.0.0.0/8")])
}
func TestBouncer_ApplyDeleted_IP(t *testing.T) {
b := newTestBouncer()
b.ips[netip.MustParseAddr("1.2.3.4")] = &restrict.CrowdSecDecision{Type: restrict.DecisionBan}
b.ips[netip.MustParseAddr("5.6.7.8")] = &restrict.CrowdSecDecision{Type: restrict.DecisionBan}
b.applyDeleted(makeDecisions(
decision{scope: "ip", value: "1.2.3.4", dtype: "ban"},
))
assert.Len(t, b.ips, 1)
assert.Nil(t, b.ips[netip.MustParseAddr("1.2.3.4")])
assert.NotNil(t, b.ips[netip.MustParseAddr("5.6.7.8")])
}
func TestBouncer_ApplyDeleted_Range(t *testing.T) {
b := newTestBouncer()
b.prefixes[netip.MustParsePrefix("10.0.0.0/8")] = &restrict.CrowdSecDecision{Type: restrict.DecisionBan}
b.prefixes[netip.MustParsePrefix("192.168.0.0/16")] = &restrict.CrowdSecDecision{Type: restrict.DecisionBan}
b.applyDeleted(makeDecisions(
decision{scope: "range", value: "10.0.0.0/8", dtype: "ban"},
))
require.Len(t, b.prefixes, 1)
assert.NotNil(t, b.prefixes[netip.MustParsePrefix("192.168.0.0/16")])
}
func TestBouncer_ApplyNew_OverwritesExisting(t *testing.T) {
b := newTestBouncer()
b.ips[netip.MustParseAddr("1.2.3.4")] = &restrict.CrowdSecDecision{Type: restrict.DecisionBan}
b.applyNew(makeDecisions(
decision{scope: "ip", value: "1.2.3.4", dtype: "captcha"},
))
assert.Equal(t, restrict.DecisionCaptcha, b.ips[netip.MustParseAddr("1.2.3.4")].Type)
}
func TestBouncer_ApplyNew_SkipsInvalid(t *testing.T) {
b := newTestBouncer()
b.applyNew(makeDecisions(
decision{scope: "ip", value: "not-an-ip", dtype: "ban"},
decision{scope: "range", value: "also-not-valid", dtype: "ban"},
))
assert.Empty(t, b.ips)
assert.Empty(t, b.prefixes)
}
// TestBouncer_StreamIntegration tests the full flow: fake LAPI → StreamBouncer → Bouncer cache → CheckIP.
func TestBouncer_StreamIntegration(t *testing.T) {
lapi := newFakeLAPI()
ts := httptest.NewServer(lapi)
defer ts.Close()
// Seed the LAPI with initial decisions.
lapi.setDecisions(
decision{scope: "ip", value: "1.2.3.4", dtype: "ban", scenario: "crowdsecurity/ssh-bf"},
decision{scope: "range", value: "10.0.0.0/8", dtype: "ban", scenario: "crowdsecurity/http-probing"},
decision{scope: "ip", value: "5.5.5.5", dtype: "captcha", scenario: "crowdsecurity/http-crawl"},
)
b := NewBouncer(ts.URL, "test-key", log.NewEntry(log.StandardLogger()))
b.tickerInterval = 200 * time.Millisecond
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
require.NoError(t, b.Start(ctx))
defer b.Stop()
// Wait for initial sync.
require.Eventually(t, b.Ready, 5*time.Second, 50*time.Millisecond, "bouncer should become ready")
// Verify decisions are cached.
d := b.CheckIP(netip.MustParseAddr("1.2.3.4"))
require.NotNil(t, d, "1.2.3.4 should be banned")
assert.Equal(t, restrict.DecisionBan, d.Type)
d2 := b.CheckIP(netip.MustParseAddr("10.1.2.3"))
require.NotNil(t, d2, "10.1.2.3 should match range ban")
assert.Equal(t, restrict.DecisionBan, d2.Type)
d3 := b.CheckIP(netip.MustParseAddr("5.5.5.5"))
require.NotNil(t, d3, "5.5.5.5 should have captcha")
assert.Equal(t, restrict.DecisionCaptcha, d3.Type)
assert.Nil(t, b.CheckIP(netip.MustParseAddr("9.9.9.9")), "unknown IP should be nil")
// Simulate a delta update: delete one IP, add a new one.
lapi.setDelta(
[]decision{{scope: "ip", value: "1.2.3.4", dtype: "ban"}},
[]decision{{scope: "ip", value: "2.3.4.5", dtype: "throttle", scenario: "crowdsecurity/http-flood"}},
)
// Wait for the delta to be picked up.
require.Eventually(t, func() bool {
return b.CheckIP(netip.MustParseAddr("2.3.4.5")) != nil
}, 5*time.Second, 50*time.Millisecond, "new decision should appear")
assert.Nil(t, b.CheckIP(netip.MustParseAddr("1.2.3.4")), "deleted decision should be gone")
d4 := b.CheckIP(netip.MustParseAddr("2.3.4.5"))
require.NotNil(t, d4)
assert.Equal(t, restrict.DecisionThrottle, d4.Type)
// Range ban should still be active.
assert.NotNil(t, b.CheckIP(netip.MustParseAddr("10.99.99.99")))
}
// Helpers
func newTestBouncer() *Bouncer {
return &Bouncer{
ips: make(map[netip.Addr]*restrict.CrowdSecDecision),
prefixes: make(map[netip.Prefix]*restrict.CrowdSecDecision),
logger: log.NewEntry(log.StandardLogger()),
}
}
type decision struct {
scope string
value string
dtype string
scenario string
}
func makeDecisions(decs ...decision) []*models.Decision {
out := make([]*models.Decision, len(decs))
for i, d := range decs {
out[i] = &models.Decision{
Scope: strPtr(d.scope),
Value: strPtr(d.value),
Type: strPtr(d.dtype),
Scenario: strPtr(d.scenario),
Duration: strPtr("1h"),
Origin: strPtr("cscli"),
}
}
return out
}
func strPtr(s string) *string { return &s }
// fakeLAPI is a minimal fake CrowdSec LAPI that serves /v1/decisions/stream.
type fakeLAPI struct {
mu sync.Mutex
initial []decision
newDelta []decision
delDelta []decision
served bool // true after the initial snapshot has been served
}
func newFakeLAPI() *fakeLAPI {
return &fakeLAPI{}
}
func (f *fakeLAPI) setDecisions(decs ...decision) {
f.mu.Lock()
defer f.mu.Unlock()
f.initial = decs
f.served = false
}
func (f *fakeLAPI) setDelta(deleted, added []decision) {
f.mu.Lock()
defer f.mu.Unlock()
f.delDelta = deleted
f.newDelta = added
}
func (f *fakeLAPI) ServeHTTP(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/v1/decisions/stream" {
http.NotFound(w, r)
return
}
f.mu.Lock()
defer f.mu.Unlock()
resp := streamResponse{}
if !f.served {
for _, d := range f.initial {
resp.New = append(resp.New, toLAPIDecision(d))
}
f.served = true
} else {
for _, d := range f.delDelta {
resp.Deleted = append(resp.Deleted, toLAPIDecision(d))
}
for _, d := range f.newDelta {
resp.New = append(resp.New, toLAPIDecision(d))
}
// Clear delta after serving once.
f.delDelta = nil
f.newDelta = nil
}
w.Header().Set("Content-Type", "application/json")
json.NewEncoder(w).Encode(resp) //nolint:errcheck
}
// streamResponse mirrors the CrowdSec LAPI /v1/decisions/stream JSON structure.
type streamResponse struct {
New []*lapiDecision `json:"new"`
Deleted []*lapiDecision `json:"deleted"`
}
type lapiDecision struct {
Duration *string `json:"duration"`
Origin *string `json:"origin"`
Scenario *string `json:"scenario"`
Scope *string `json:"scope"`
Type *string `json:"type"`
Value *string `json:"value"`
}
func toLAPIDecision(d decision) *lapiDecision {
return &lapiDecision{
Duration: strPtr("1h"),
Origin: strPtr("cscli"),
Scenario: strPtr(d.scenario),
Scope: strPtr(d.scope),
Type: strPtr(d.dtype),
Value: strPtr(d.value),
}
}
proxy/internal/crowdsec/registry.go
+103
View File
@@ -0,0 +1,103 @@
package crowdsec
import (
"context"
"sync"
log "github.com/sirupsen/logrus"
"github.com/netbirdio/netbird/proxy/internal/types"
)
// Registry manages a single shared Bouncer instance with reference counting.
// The bouncer starts when the first service acquires it and stops when the
// last service releases it.
type Registry struct {
mu sync.Mutex
bouncer *Bouncer
refs map[types.ServiceID]struct{}
apiURL string
apiKey string
logger *log.Entry
cancel context.CancelFunc
}
// NewRegistry creates a registry. The bouncer is not started until Acquire is called.
func NewRegistry(apiURL, apiKey string, logger *log.Entry) *Registry {
return &Registry{
apiURL: apiURL,
apiKey: apiKey,
logger: logger,
refs: make(map[types.ServiceID]struct{}),
}
}
// Available returns true when the LAPI URL and API key are configured.
func (r *Registry) Available() bool {
return r.apiURL != "" && r.apiKey != ""
}
// Acquire registers svcID as a consumer and starts the bouncer if this is the
// first consumer. Returns the shared Bouncer (which implements the restrict
// package's CrowdSecChecker interface). Returns nil if not Available.
func (r *Registry) Acquire(svcID types.ServiceID) *Bouncer {
r.mu.Lock()
defer r.mu.Unlock()
if !r.Available() {
return nil
}
if _, exists := r.refs[svcID]; exists {
return r.bouncer
}
if r.bouncer == nil {
r.startLocked()
}
// startLocked may fail, leaving r.bouncer nil.
if r.bouncer == nil {
return nil
}
r.refs[svcID] = struct{}{}
return r.bouncer
}
// Release removes svcID as a consumer. Stops the bouncer when the last
// consumer releases.
func (r *Registry) Release(svcID types.ServiceID) {
r.mu.Lock()
defer r.mu.Unlock()
delete(r.refs, svcID)
if len(r.refs) == 0 && r.bouncer != nil {
r.stopLocked()
}
}
func (r *Registry) startLocked() {
b := NewBouncer(r.apiURL, r.apiKey, r.logger)
ctx, cancel := context.WithCancel(context.Background())
r.cancel = cancel
if err := b.Start(ctx); err != nil {
r.logger.Errorf("failed to start CrowdSec bouncer: %v", err)
cancel()
return
}
r.bouncer = b
r.logger.Info("CrowdSec bouncer started")
}
func (r *Registry) stopLocked() {
r.bouncer.Stop()
r.cancel()
r.bouncer = nil
r.cancel = nil
r.logger.Info("CrowdSec bouncer stopped")
}
proxy/internal/crowdsec/registry_test.go
+66
View File
@@ -0,0 +1,66 @@
package crowdsec
import (
"testing"
log "github.com/sirupsen/logrus"
"github.com/stretchr/testify/assert"
"github.com/netbirdio/netbird/proxy/internal/types"
)
func TestRegistry_Available(t *testing.T) {
r := NewRegistry("http://localhost:8080/", "test-key", log.NewEntry(log.StandardLogger()))
assert.True(t, r.Available())
r2 := NewRegistry("", "", log.NewEntry(log.StandardLogger()))
assert.False(t, r2.Available())
r3 := NewRegistry("http://localhost:8080/", "", log.NewEntry(log.StandardLogger()))
assert.False(t, r3.Available())
}
func TestRegistry_Acquire_NotAvailable(t *testing.T) {
r := NewRegistry("", "", log.NewEntry(log.StandardLogger()))
b := r.Acquire("svc-1")
assert.Nil(t, b)
}
func TestRegistry_Acquire_Idempotent(t *testing.T) {
r := newTestRegistry()
b1 := r.Acquire("svc-1")
// Can't start without a real LAPI, but we can verify the ref tracking.
// The bouncer will be nil because Start fails, but the ref is tracked.
_ = b1
assert.Len(t, r.refs, 1)
// Second acquire of same service should not add another ref.
r.Acquire("svc-1")
assert.Len(t, r.refs, 1)
}
func TestRegistry_Release_Removes(t *testing.T) {
r := newTestRegistry()
r.refs[types.ServiceID("svc-1")] = struct{}{}
r.Release("svc-1")
assert.Empty(t, r.refs)
}
func TestRegistry_Release_Noop(t *testing.T) {
r := newTestRegistry()
// Releasing a service that was never acquired should not panic.
r.Release("nonexistent")
assert.Empty(t, r.refs)
}
func newTestRegistry() *Registry {
return &Registry{
apiURL: "http://localhost:8080/",
apiKey: "test-key",
logger: log.NewEntry(log.StandardLogger()),
refs: make(map[types.ServiceID]struct{}),
}
}
proxy/internal/debug/client.go
+576
View File
@@ -0,0 +1,576 @@
// Package debug provides HTTP debug endpoints and CLI client for the proxy server.
package debug
import (
"context"
"encoding/json"
"fmt"
"io"
"net"
"net/http"
"net/url"
"strings"
"time"
)
// StatusFilters contains filter options for status queries.
type StatusFilters struct {
IPs []string
Names []string
Status string
ConnectionType string
}
// Client provides CLI access to debug endpoints.
type Client struct {
baseURL string
jsonOutput bool
httpClient *http.Client
out io.Writer
}
// NewClient creates a new debug client.
func NewClient(baseURL string, jsonOutput bool, out io.Writer) *Client {
if !strings.HasPrefix(baseURL, "http://") && !strings.HasPrefix(baseURL, "https://") {
baseURL = "http://" + baseURL
}
baseURL = strings.TrimSuffix(baseURL, "/")
return &Client{
baseURL: baseURL,
jsonOutput: jsonOutput,
out: out,
httpClient: &http.Client{
Timeout: 30 * time.Second,
},
}
}
// Health fetches the health status.
func (c *Client) Health(ctx context.Context) error {
return c.fetchAndPrint(ctx, "/debug/health", c.printHealth)
}
func (c *Client) printHealth(data map[string]any) {
_, _ = fmt.Fprintf(c.out, "Status: %v\n", data["status"])
_, _ = fmt.Fprintf(c.out, "Uptime: %v\n", data["uptime"])
_, _ = fmt.Fprintf(c.out, "Management Connected: %s\n", boolIcon(data["management_connected"]))
_, _ = fmt.Fprintf(c.out, "All Clients Healthy: %s\n", boolIcon(data["all_clients_healthy"]))
total, _ := data["certs_total"].(float64)
ready, _ := data["certs_ready"].(float64)
pending, _ := data["certs_pending"].(float64)
failed, _ := data["certs_failed"].(float64)
if total > 0 {
_, _ = fmt.Fprintf(c.out, "Certificates: %d ready, %d pending, %d failed (%d total)\n",
int(ready), int(pending), int(failed), int(total))
}
if domains, ok := data["certs_ready_domains"].([]any); ok && len(domains) > 0 {
_, _ = fmt.Fprintf(c.out, " Ready:\n")
for _, d := range domains {
_, _ = fmt.Fprintf(c.out, " %v\n", d)
}
}
if domains, ok := data["certs_pending_domains"].([]any); ok && len(domains) > 0 {
_, _ = fmt.Fprintf(c.out, " Pending:\n")
for _, d := range domains {
_, _ = fmt.Fprintf(c.out, " %v\n", d)
}
}
if domains, ok := data["certs_failed_domains"].(map[string]any); ok && len(domains) > 0 {
_, _ = fmt.Fprintf(c.out, " Failed:\n")
for d, errMsg := range domains {
_, _ = fmt.Fprintf(c.out, " %s: %v\n", d, errMsg)
}
}
c.printHealthClients(data)
}
func (c *Client) printHealthClients(data map[string]any) {
clients, ok := data["clients"].(map[string]any)
if !ok || len(clients) == 0 {
return
}
_, _ = fmt.Fprintf(c.out, "\n%-38s %-9s %-7s %-8s %-8s %-16s %s\n",
"ACCOUNT ID", "HEALTHY", "MGMT", "SIGNAL", "RELAYS", "PEERS (P2P/RLY)", "DEGRADED")
_, _ = fmt.Fprintln(c.out, strings.Repeat("-", 110))
for accountID, v := range clients {
ch, ok := v.(map[string]any)
if !ok {
continue
}
healthy := boolIcon(ch["healthy"])
mgmt := boolIcon(ch["management_connected"])
signal := boolIcon(ch["signal_connected"])
relaysConn, _ := ch["relays_connected"].(float64)
relaysTotal, _ := ch["relays_total"].(float64)
relays := fmt.Sprintf("%d/%d", int(relaysConn), int(relaysTotal))
peersConnected, _ := ch["peers_connected"].(float64)
peersTotal, _ := ch["peers_total"].(float64)
peersP2P, _ := ch["peers_p2p"].(float64)
peersRelayed, _ := ch["peers_relayed"].(float64)
peersDegraded, _ := ch["peers_degraded"].(float64)
peers := fmt.Sprintf("%d/%d (%d/%d)", int(peersConnected), int(peersTotal), int(peersP2P), int(peersRelayed))
degraded := fmt.Sprintf("%d", int(peersDegraded))
_, _ = fmt.Fprintf(c.out, "%-38s %-9s %-7s %-8s %-8s %-16s %s", accountID, healthy, mgmt, signal, relays, peers, degraded)
if errMsg, ok := ch["error"].(string); ok && errMsg != "" {
_, _ = fmt.Fprintf(c.out, " (%s)", errMsg)
}
_, _ = fmt.Fprintln(c.out)
}
}
func boolIcon(v any) string {
b, ok := v.(bool)
if !ok {
return "?"
}
if b {
return "yes"
}
return "no"
}
// ListClients fetches the list of all clients.
func (c *Client) ListClients(ctx context.Context) error {
return c.fetchAndPrint(ctx, "/debug/clients", c.printClients)
}
func (c *Client) printClients(data map[string]any) {
_, _ = fmt.Fprintf(c.out, "Uptime: %v\n", data["uptime"])
_, _ = fmt.Fprintf(c.out, "Clients: %v\n\n", data["client_count"])
clients, ok := data["clients"].([]any)
if !ok || len(clients) == 0 {
_, _ = fmt.Fprintln(c.out, "No clients connected.")
return
}
_, _ = fmt.Fprintf(c.out, "%-38s %-12s %-40s %s\n", "ACCOUNT ID", "AGE", "SERVICES", "HAS CLIENT")
_, _ = fmt.Fprintln(c.out, strings.Repeat("-", 110))
for _, item := range clients {
c.printClientRow(item)
}
c.printInboundListeners(clients)
}
func (c *Client) printInboundListeners(clients []any) {
type row struct {
accountID string
tunnelIP string
httpsPort int
httpPort int
}
var rows []row
for _, item := range clients {
client, ok := item.(map[string]any)
if !ok {
continue
}
inbound, ok := client["inbound_listener"].(map[string]any)
if !ok {
continue
}
tunnelIP, _ := inbound["tunnel_ip"].(string)
httpsPort, _ := inbound["https_port"].(float64)
httpPort, _ := inbound["http_port"].(float64)
accountID, _ := client["account_id"].(string)
rows = append(rows, row{
accountID: accountID,
tunnelIP: tunnelIP,
httpsPort: int(httpsPort),
httpPort: int(httpPort),
})
}
if len(rows) == 0 {
return
}
_, _ = fmt.Fprintln(c.out)
_, _ = fmt.Fprintln(c.out, "Inbound listeners (per-account):")
_, _ = fmt.Fprintf(c.out, " %-38s %-20s %-7s %s\n", "ACCOUNT ID", "TUNNEL IP", "HTTPS", "HTTP")
_, _ = fmt.Fprintln(c.out, " "+strings.Repeat("-", 78))
for _, r := range rows {
_, _ = fmt.Fprintf(c.out, " %-38s %-20s %-7d %d\n", r.accountID, r.tunnelIP, r.httpsPort, r.httpPort)
}
}
func (c *Client) printClientRow(item any) {
client, ok := item.(map[string]any)
if !ok {
return
}
services := c.extractServiceKeys(client)
hasClient := "no"
if hc, ok := client["has_client"].(bool); ok && hc {
hasClient = "yes"
}
_, _ = fmt.Fprintf(c.out, "%-38s %-12v %s %s\n",
client["account_id"],
client["age"],
services,
hasClient,
)
}
func (c *Client) extractServiceKeys(client map[string]any) string {
d, ok := client["service_keys"].([]any)
if !ok || len(d) == 0 {
return "-"
}
parts := make([]string, len(d))
for i, key := range d {
parts[i] = fmt.Sprint(key)
}
return strings.Join(parts, ", ")
}
// ClientStatus fetches the status of a specific client.
func (c *Client) ClientStatus(ctx context.Context, accountID string, filters StatusFilters) error {
params := url.Values{}
if len(filters.IPs) > 0 {
params.Set("filter-by-ips", strings.Join(filters.IPs, ","))
}
if len(filters.Names) > 0 {
params.Set("filter-by-names", strings.Join(filters.Names, ","))
}
if filters.Status != "" {
params.Set("filter-by-status", filters.Status)
}
if filters.ConnectionType != "" {
params.Set("filter-by-connection-type", filters.ConnectionType)
}
path := "/debug/clients/" + url.PathEscape(accountID)
if len(params) > 0 {
path += "?" + params.Encode()
}
return c.fetchAndPrint(ctx, path, c.printClientStatus)
}
func (c *Client) printClientStatus(data map[string]any) {
_, _ = fmt.Fprintf(c.out, "Account: %v\n", data["account_id"])
if inbound, ok := data["inbound_listener"].(map[string]any); ok {
tunnelIP, _ := inbound["tunnel_ip"].(string)
httpsPort, _ := inbound["https_port"].(float64)
httpPort, _ := inbound["http_port"].(float64)
_, _ = fmt.Fprintf(c.out, "Inbound listener: %s (https=%d, http=%d)\n", tunnelIP, int(httpsPort), int(httpPort))
}
_, _ = fmt.Fprintln(c.out)
if status, ok := data["status"].(string); ok {
_, _ = fmt.Fprint(c.out, status)
}
}
// ClientSyncResponse fetches the sync response of a specific client.
func (c *Client) ClientSyncResponse(ctx context.Context, accountID string) error {
path := "/debug/clients/" + url.PathEscape(accountID) + "/syncresponse"
return c.fetchAndPrintJSON(ctx, path)
}
// PingTCP performs a TCP ping through a client.
// ipVersion may be "4", "6", or "" for automatic.
func (c *Client) PingTCP(ctx context.Context, accountID, host string, port int, timeout time.Duration, ipVersion string) error {
params := url.Values{}
params.Set("host", host)
params.Set("port", fmt.Sprintf("%d", port))
if timeout > 0 {
params.Set("timeout", timeout.String())
}
if ipVersion != "" {
params.Set("ip_version", ipVersion)
}
path := fmt.Sprintf("/debug/clients/%s/pingtcp?%s", url.PathEscape(accountID), params.Encode())
return c.fetchAndPrint(ctx, path, c.printPingResult)
}
func (c *Client) printPingResult(data map[string]any) {
success, _ := data["success"].(bool)
host := net.JoinHostPort(fmt.Sprint(data["host"]), fmt.Sprint(data["port"]))
if success {
remote, _ := data["remote"].(string)
if remote != "" && remote != host {
_, _ = fmt.Fprintf(c.out, "Success: %s (via %s)\n", host, remote)
} else {
_, _ = fmt.Fprintf(c.out, "Success: %s\n", host)
}
_, _ = fmt.Fprintf(c.out, "Latency: %v\n", data["latency"])
} else {
_, _ = fmt.Fprintf(c.out, "Failed: %s\n", host)
c.printError(data)
}
}
// SetLogLevel sets the log level of a specific client.
func (c *Client) SetLogLevel(ctx context.Context, accountID, level string) error {
params := url.Values{}
params.Set("level", level)
path := fmt.Sprintf("/debug/clients/%s/loglevel?%s", url.PathEscape(accountID), params.Encode())
return c.fetchAndPrint(ctx, path, c.printLogLevelResult)
}
func (c *Client) printLogLevelResult(data map[string]any) {
success, _ := data["success"].(bool)
if success {
_, _ = fmt.Fprintf(c.out, "Log level set to: %v\n", data["level"])
} else {
_, _ = fmt.Fprintln(c.out, "Failed to set log level")
c.printError(data)
}
}
// PerfSet live-retunes the tunnel buffer pool cap on all running embedded
// clients. Batch size is not live-tunable; configure it at proxy startup.
func (c *Client) PerfSet(ctx context.Context, value uint32) error {
path := fmt.Sprintf("/debug/perf?value=%d", value)
return c.fetchAndPrint(ctx, path, c.printPerfSet)
}
func (c *Client) printPerfSet(data map[string]any) {
if errMsg, ok := data["error"].(string); ok && errMsg != "" {
c.printError(data)
return
}
val, _ := data["value"].(float64)
applied, _ := data["applied"].(float64)
_, _ = fmt.Fprintf(c.out, "Pool cap set to: %d\n", uint32(val))
_, _ = fmt.Fprintf(c.out, "Applied to %d live clients\n", int(applied))
if failed, ok := data["failed"].(map[string]any); ok && len(failed) > 0 {
_, _ = fmt.Fprintln(c.out, "Failed:")
for k, v := range failed {
_, _ = fmt.Fprintf(c.out, " %s: %v\n", k, v)
}
}
}
// Runtime fetches runtime stats (heap, goroutines, RSS).
func (c *Client) Runtime(ctx context.Context) error {
return c.fetchAndPrint(ctx, "/debug/runtime", c.printRuntime)
}
func (c *Client) printRuntime(data map[string]any) {
i := func(k string) uint64 {
v, _ := data[k].(float64)
return uint64(v)
}
mb := func(n uint64) string { return fmt.Sprintf("%.1f MB", float64(n)/(1<<20)) }
_, _ = fmt.Fprintf(c.out, "Uptime: %v\n", data["uptime"])
_, _ = fmt.Fprintf(c.out, "Go: %v on %d CPU (GOMAXPROCS=%d)\n", data["go_version"], uint32(i("num_cpu")), uint32(i("gomaxprocs")))
_, _ = fmt.Fprintf(c.out, "Goroutines: %d\n", i("goroutines"))
_, _ = fmt.Fprintf(c.out, "Live objects: %d\n", i("live_objects"))
_, _ = fmt.Fprintf(c.out, "GC: %d cycles, %v pause total\n", i("num_gc"), time.Duration(i("pause_total_ns")))
_, _ = fmt.Fprintln(c.out, "Heap:")
_, _ = fmt.Fprintf(c.out, " alloc: %s\n", mb(i("heap_alloc")))
_, _ = fmt.Fprintf(c.out, " in-use: %s\n", mb(i("heap_inuse")))
_, _ = fmt.Fprintf(c.out, " idle: %s\n", mb(i("heap_idle")))
_, _ = fmt.Fprintf(c.out, " released: %s\n", mb(i("heap_released")))
_, _ = fmt.Fprintf(c.out, " sys: %s\n", mb(i("heap_sys")))
_, _ = fmt.Fprintf(c.out, "Total sys: %s\n", mb(i("sys")))
if _, ok := data["vm_rss"]; ok {
_, _ = fmt.Fprintln(c.out, "Process:")
_, _ = fmt.Fprintf(c.out, " VmRSS: %s\n", mb(i("vm_rss")))
_, _ = fmt.Fprintf(c.out, " VmSize: %s\n", mb(i("vm_size")))
_, _ = fmt.Fprintf(c.out, " VmData: %s\n", mb(i("vm_data")))
}
_, _ = fmt.Fprintf(c.out, "Clients: %d (%d started)\n", i("clients"), i("started"))
}
// StartClient starts a specific client.
func (c *Client) StartClient(ctx context.Context, accountID string) error {
path := "/debug/clients/" + url.PathEscape(accountID) + "/start"
return c.fetchAndPrint(ctx, path, c.printStartResult)
}
func (c *Client) printStartResult(data map[string]any) {
success, _ := data["success"].(bool)
if success {
_, _ = fmt.Fprintln(c.out, "Client started")
} else {
_, _ = fmt.Fprintln(c.out, "Failed to start client")
c.printError(data)
}
}
// StopClient stops a specific client.
func (c *Client) StopClient(ctx context.Context, accountID string) error {
path := "/debug/clients/" + url.PathEscape(accountID) + "/stop"
return c.fetchAndPrint(ctx, path, c.printStopResult)
}
func (c *Client) printStopResult(data map[string]any) {
success, _ := data["success"].(bool)
if success {
_, _ = fmt.Fprintln(c.out, "Client stopped")
} else {
_, _ = fmt.Fprintln(c.out, "Failed to stop client")
c.printError(data)
}
}
func (c *Client) printError(data map[string]any) {
if errMsg, ok := data["error"].(string); ok {
_, _ = fmt.Fprintf(c.out, "Error: %s\n", errMsg)
}
}
// CaptureOptions configures a capture request.
type CaptureOptions struct {
AccountID string
Duration string
FilterExpr string
Text bool
Verbose bool
ASCII bool
Output io.Writer
}
// Capture streams a packet capture from the debug endpoint. The response body
// (pcap or text) is written directly to opts.Output until the server closes the
// connection or the context is cancelled.
func (c *Client) Capture(ctx context.Context, opts CaptureOptions) error {
if opts.AccountID == "" {
return fmt.Errorf("account ID is required")
}
if opts.Output == nil {
return fmt.Errorf("output writer is required")
}
params := url.Values{}
if opts.Duration != "" {
params.Set("duration", opts.Duration)
}
if opts.FilterExpr != "" {
params.Set("filter", opts.FilterExpr)
}
if opts.Text {
params.Set("format", "text")
}
if opts.Verbose {
params.Set("verbose", "true")
}
if opts.ASCII {
params.Set("ascii", "true")
}
path := fmt.Sprintf("/debug/clients/%s/capture", url.PathEscape(opts.AccountID))
if len(params) > 0 {
path += "?" + params.Encode()
}
fullURL := c.baseURL + path
req, err := http.NewRequestWithContext(ctx, http.MethodGet, fullURL, nil)
if err != nil {
return fmt.Errorf("create request: %w", err)
}
// Use a separate client without timeout since captures stream for their full duration.
httpClient := &http.Client{}
resp, err := httpClient.Do(req)
if err != nil {
return fmt.Errorf("request failed: %w", err)
}
defer func() { _ = resp.Body.Close() }()
if resp.StatusCode >= 400 {
body, _ := io.ReadAll(resp.Body)
return fmt.Errorf("server error (%d): %s", resp.StatusCode, strings.TrimSpace(string(body)))
}
_, err = io.Copy(opts.Output, resp.Body)
if err != nil && ctx.Err() != nil {
return nil
}
return err
}
func (c *Client) fetchAndPrint(ctx context.Context, path string, printer func(map[string]any)) error {
data, raw, err := c.fetch(ctx, path)
if err != nil {
return err
}
if c.jsonOutput {
return c.writeJSON(data)
}
if data != nil {
printer(data)
return nil
}
_, _ = fmt.Fprintln(c.out, string(raw))
return nil
}
func (c *Client) fetchAndPrintJSON(ctx context.Context, path string) error {
data, raw, err := c.fetch(ctx, path)
if err != nil {
return err
}
if data != nil {
return c.writeJSON(data)
}
_, _ = fmt.Fprintln(c.out, string(raw))
return nil
}
func (c *Client) writeJSON(data map[string]any) error {
enc := json.NewEncoder(c.out)
enc.SetIndent("", " ")
return enc.Encode(data)
}
func (c *Client) fetch(ctx context.Context, path string) (map[string]any, []byte, error) {
fullURL := c.baseURL + path
if !strings.Contains(path, "format=json") {
if strings.Contains(path, "?") {
fullURL += "&format=json"
} else {
fullURL += "?format=json"
}
}
req, err := http.NewRequestWithContext(ctx, http.MethodGet, fullURL, nil)
if err != nil {
return nil, nil, fmt.Errorf("create request: %w", err)
}
resp, err := c.httpClient.Do(req)
if err != nil {
return nil, nil, fmt.Errorf("request failed: %w", err)
}
defer func() { _ = resp.Body.Close() }()
body, err := io.ReadAll(resp.Body)
if err != nil {
return nil, nil, fmt.Errorf("read response: %w", err)
}
if resp.StatusCode >= 400 {
return nil, nil, fmt.Errorf("server error (%d): %s", resp.StatusCode, strings.TrimSpace(string(body)))
}
var data map[string]any
if err := json.Unmarshal(body, &data); err != nil {
return nil, body, nil
}
return data, body, nil
}
proxy/internal/debug/client_test.go
+71
View File
@@ -0,0 +1,71 @@
package debug
import (
"bytes"
"testing"
"github.com/stretchr/testify/assert"
)
func TestPrintHealth_WithCertsAndClients(t *testing.T) {
var buf bytes.Buffer
c := NewClient("localhost:8444", false, &buf)
data := map[string]any{
"status": "ok",
"uptime": "1h30m",
"management_connected": true,
"all_clients_healthy": true,
"certs_total": float64(3),
"certs_ready": float64(2),
"certs_pending": float64(1),
"certs_failed": float64(0),
"certs_ready_domains": []any{"a.example.com", "b.example.com"},
"certs_pending_domains": []any{"c.example.com"},
"clients": map[string]any{
"acc-1": map[string]any{
"healthy": true,
"management_connected": true,
"signal_connected": true,
"relays_connected": float64(1),
"relays_total": float64(2),
"peers_connected": float64(3),
"peers_total": float64(5),
"peers_p2p": float64(2),
"peers_relayed": float64(1),
"peers_degraded": float64(0),
},
},
}
c.printHealth(data)
out := buf.String()
assert.Contains(t, out, "Status: ok")
assert.Contains(t, out, "Uptime: 1h30m")
assert.Contains(t, out, "yes") // management_connected
assert.Contains(t, out, "2 ready, 1 pending, 0 failed (3 total)")
assert.Contains(t, out, "a.example.com")
assert.Contains(t, out, "c.example.com")
assert.Contains(t, out, "acc-1")
}
func TestPrintHealth_Minimal(t *testing.T) {
var buf bytes.Buffer
c := NewClient("localhost:8444", false, &buf)
data := map[string]any{
"status": "ok",
"uptime": "5m",
"management_connected": false,
"all_clients_healthy": false,
}
c.printHealth(data)
out := buf.String()
assert.Contains(t, out, "Status: ok")
assert.Contains(t, out, "Uptime: 5m")
assert.NotContains(t, out, "Certificates")
assert.NotContains(t, out, "ACCOUNT ID")
}
proxy/internal/debug/handler.go
+960
View File
@@ -0,0 +1,960 @@
// Package debug provides HTTP debug endpoints for the proxy server.
package debug
import (
"cmp"
"context"
"embed"
"encoding/json"
"fmt"
"html/template"
"maps"
"net"
"net/http"
"os"
"runtime"
"slices"
"strconv"
"strings"
"sync"
"time"
log "github.com/sirupsen/logrus"
"google.golang.org/protobuf/encoding/protojson"
nbembed "github.com/netbirdio/netbird/client/embed"
nbstatus "github.com/netbirdio/netbird/client/status"
"github.com/netbirdio/netbird/proxy/internal/health"
"github.com/netbirdio/netbird/proxy/internal/roundtrip"
"github.com/netbirdio/netbird/proxy/internal/types"
"github.com/netbirdio/netbird/version"
)
//go:embed templates/*.html
var templateFS embed.FS
const defaultPingTimeout = 10 * time.Second
// formatDuration formats a duration with 2 decimal places using appropriate units.
func formatDuration(d time.Duration) string {
switch {
case d >= time.Hour:
return fmt.Sprintf("%.2fh", d.Hours())
case d >= time.Minute:
return fmt.Sprintf("%.2fm", d.Minutes())
case d >= time.Second:
return fmt.Sprintf("%.2fs", d.Seconds())
case d >= time.Millisecond:
return fmt.Sprintf("%.2fms", float64(d.Microseconds())/1000)
case d >= time.Microsecond:
return fmt.Sprintf("%.2fµs", float64(d.Nanoseconds())/1000)
default:
return fmt.Sprintf("%dns", d.Nanoseconds())
}
}
func sortedAccountIDs(m map[types.AccountID]roundtrip.ClientDebugInfo) []types.AccountID {
return slices.Sorted(maps.Keys(m))
}
// clientProvider provides access to NetBird clients.
type clientProvider interface {
GetClient(accountID types.AccountID) (*nbembed.Client, bool)
ListClientsForDebug() map[types.AccountID]roundtrip.ClientDebugInfo
ListClientsForStartup() map[types.AccountID]*nbembed.Client
}
// InboundListenerInfo describes a per-account inbound listener as
// surfaced through the debug HTTP handler. Mirrors the proto sub-message
// emitted with SendStatusUpdate so dashboards and CLI tooling see the
// same shape.
type InboundListenerInfo struct {
TunnelIP string `json:"tunnel_ip"`
HTTPSPort uint16 `json:"https_port"`
HTTPPort uint16 `json:"http_port"`
}
// InboundProvider exposes per-account inbound listener state. Optional;
// when nil the debug endpoint omits the inbound section entirely so the
// existing JSON shape stays additive.
type InboundProvider interface {
InboundListeners() map[types.AccountID]InboundListenerInfo
}
// healthChecker provides health probe state.
type healthChecker interface {
ReadinessProbe() bool
StartupProbe(ctx context.Context) bool
CheckClientsConnected(ctx context.Context) (bool, map[types.AccountID]health.ClientHealth)
}
type certStatus interface {
TotalDomains() int
PendingDomains() []string
ReadyDomains() []string
FailedDomains() map[string]string
}
// Handler provides HTTP debug endpoints.
type Handler struct {
provider clientProvider
health healthChecker
certStatus certStatus
inbound InboundProvider
logger *log.Logger
startTime time.Time
templates *template.Template
templateMu sync.RWMutex
}
// NewHandler creates a new debug handler.
func NewHandler(provider clientProvider, healthChecker healthChecker, logger *log.Logger) *Handler {
if logger == nil {
logger = log.StandardLogger()
}
h := &Handler{
provider: provider,
health: healthChecker,
logger: logger,
startTime: time.Now(),
}
if err := h.loadTemplates(); err != nil {
logger.Errorf("failed to load embedded templates: %v", err)
}
return h
}
// SetCertStatus sets the certificate status provider for ACME prefetch observability.
func (h *Handler) SetCertStatus(cs certStatus) {
h.certStatus = cs
}
// SetInboundProvider wires per-account inbound listener observability.
// Pass nil (or skip the call) to keep the inbound section out of debug
// responses on proxies that don't run --private.
func (h *Handler) SetInboundProvider(p InboundProvider) {
h.inbound = p
}
func (h *Handler) loadTemplates() error {
tmpl, err := template.ParseFS(templateFS, "templates/*.html")
if err != nil {
return fmt.Errorf("parse embedded templates: %w", err)
}
h.templateMu.Lock()
h.templates = tmpl
h.templateMu.Unlock()
return nil
}
func (h *Handler) getTemplates() *template.Template {
h.templateMu.RLock()
defer h.templateMu.RUnlock()
return h.templates
}
// ServeHTTP handles debug requests.
func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
path := r.URL.Path
wantJSON := r.URL.Query().Get("format") == "json" || strings.HasSuffix(path, "/json")
path = strings.TrimSuffix(path, "/json")
switch path {
case "/debug", "/debug/":
h.handleIndex(w, r, wantJSON)
case "/debug/clients":
h.handleListClients(w, r, wantJSON)
case "/debug/health":
h.handleHealth(w, r, wantJSON)
case "/debug/perf":
h.handlePerf(w, r)
case "/debug/runtime":
h.handleRuntime(w, r)
default:
if h.handleClientRoutes(w, r, path, wantJSON) {
return
}
http.NotFound(w, r)
}
}
func (h *Handler) handleClientRoutes(w http.ResponseWriter, r *http.Request, path string, wantJSON bool) bool {
if !strings.HasPrefix(path, "/debug/clients/") {
return false
}
rest := strings.TrimPrefix(path, "/debug/clients/")
parts := strings.SplitN(rest, "/", 2)
accountID := types.AccountID(parts[0])
if len(parts) == 1 {
h.handleClientStatus(w, r, accountID, wantJSON)
return true
}
switch parts[1] {
case "syncresponse":
h.handleClientSyncResponse(w, r, accountID, wantJSON)
case "tools":
h.handleClientTools(w, r, accountID)
case "pingtcp":
h.handlePingTCP(w, r, accountID)
case "loglevel":
h.handleLogLevel(w, r, accountID)
case "start":
h.handleClientStart(w, r, accountID)
case "stop":
h.handleClientStop(w, r, accountID)
case "capture":
h.handleCapture(w, r, accountID)
default:
return false
}
return true
}
type failedDomain struct {
Domain string
Error string
}
type indexData struct {
Version string
Uptime string
ClientCount int
TotalServices int
CertsTotal int
CertsReady int
CertsPending int
CertsFailed int
CertsPendingDomains []string
CertsReadyDomains []string
CertsFailedDomains []failedDomain
Clients []clientData
}
type clientData struct {
AccountID string
Services string
Age string
Status string
}
func (h *Handler) handleIndex(w http.ResponseWriter, _ *http.Request, wantJSON bool) {
clients := h.provider.ListClientsForDebug()
sortedIDs := sortedAccountIDs(clients)
totalServices := 0
for _, info := range clients {
totalServices += info.ServiceCount
}
var certsTotal, certsReady, certsPending, certsFailed int
var certsPendingDomains, certsReadyDomains []string
var certsFailedDomains map[string]string
if h.certStatus != nil {
certsTotal = h.certStatus.TotalDomains()
certsPendingDomains = h.certStatus.PendingDomains()
certsReadyDomains = h.certStatus.ReadyDomains()
certsFailedDomains = h.certStatus.FailedDomains()
certsReady = len(certsReadyDomains)
certsPending = len(certsPendingDomains)
certsFailed = len(certsFailedDomains)
}
if wantJSON {
clientsJSON := make([]map[string]any, 0, len(clients))
for _, id := range sortedIDs {
info := clients[id]
clientsJSON = append(clientsJSON, map[string]any{
"account_id": info.AccountID,
"service_count": info.ServiceCount,
"service_keys": info.ServiceKeys,
"has_client": info.HasClient,
"created_at": info.CreatedAt,
"age": time.Since(info.CreatedAt).Round(time.Second).String(),
})
}
resp := map[string]any{
"version": version.NetbirdVersion(),
"uptime": time.Since(h.startTime).Round(time.Second).String(),
"client_count": len(clients),
"total_services": totalServices,
"certs_total": certsTotal,
"certs_ready": certsReady,
"certs_pending": certsPending,
"certs_failed": certsFailed,
"clients": clientsJSON,
}
if len(certsPendingDomains) > 0 {
resp["certs_pending_domains"] = certsPendingDomains
}
if len(certsReadyDomains) > 0 {
resp["certs_ready_domains"] = certsReadyDomains
}
if len(certsFailedDomains) > 0 {
resp["certs_failed_domains"] = certsFailedDomains
}
h.writeJSON(w, resp)
return
}
sortedFailed := make([]failedDomain, 0, len(certsFailedDomains))
for d, e := range certsFailedDomains {
sortedFailed = append(sortedFailed, failedDomain{Domain: d, Error: e})
}
slices.SortFunc(sortedFailed, func(a, b failedDomain) int {
return cmp.Compare(a.Domain, b.Domain)
})
data := indexData{
Version: version.NetbirdVersion(),
Uptime: time.Since(h.startTime).Round(time.Second).String(),
ClientCount: len(clients),
TotalServices: totalServices,
CertsTotal: certsTotal,
CertsReady: certsReady,
CertsPending: certsPending,
CertsFailed: certsFailed,
CertsPendingDomains: certsPendingDomains,
CertsReadyDomains: certsReadyDomains,
CertsFailedDomains: sortedFailed,
Clients: make([]clientData, 0, len(clients)),
}
for _, id := range sortedIDs {
info := clients[id]
services := strings.Join(info.ServiceKeys, ", ")
if services == "" {
services = "-"
}
status := "No client"
if info.HasClient {
status = "Active"
}
data.Clients = append(data.Clients, clientData{
AccountID: string(info.AccountID),
Services: services,
Age: time.Since(info.CreatedAt).Round(time.Second).String(),
Status: status,
})
}
h.renderTemplate(w, "index", data)
}
type clientsData struct {
Uptime string
Clients []clientData
}
func (h *Handler) handleListClients(w http.ResponseWriter, _ *http.Request, wantJSON bool) {
clients := h.provider.ListClientsForDebug()
sortedIDs := sortedAccountIDs(clients)
if wantJSON {
var inboundAll map[types.AccountID]InboundListenerInfo
if h.inbound != nil {
inboundAll = h.inbound.InboundListeners()
}
clientsJSON := make([]map[string]any, 0, len(clients))
for _, id := range sortedIDs {
info := clients[id]
row := map[string]any{
"account_id": info.AccountID,
"service_count": info.ServiceCount,
"service_keys": info.ServiceKeys,
"has_client": info.HasClient,
"created_at": info.CreatedAt,
"age": time.Since(info.CreatedAt).Round(time.Second).String(),
}
if inb, ok := inboundAll[id]; ok {
row["inbound_listener"] = inb
}
clientsJSON = append(clientsJSON, row)
}
resp := map[string]any{
"uptime": time.Since(h.startTime).Round(time.Second).String(),
"client_count": len(clients),
"clients": clientsJSON,
}
if len(inboundAll) > 0 {
resp["inbound_listener_count"] = len(inboundAll)
}
h.writeJSON(w, resp)
return
}
data := clientsData{
Uptime: time.Since(h.startTime).Round(time.Second).String(),
Clients: make([]clientData, 0, len(clients)),
}
for _, id := range sortedIDs {
info := clients[id]
services := strings.Join(info.ServiceKeys, ", ")
if services == "" {
services = "-"
}
status := "No client"
if info.HasClient {
status = "Active"
}
data.Clients = append(data.Clients, clientData{
AccountID: string(info.AccountID),
Services: services,
Age: time.Since(info.CreatedAt).Round(time.Second).String(),
Status: status,
})
}
h.renderTemplate(w, "clients", data)
}
type clientDetailData struct {
AccountID string
ActiveTab string
Content string
}
func (h *Handler) handleClientStatus(w http.ResponseWriter, r *http.Request, accountID types.AccountID, wantJSON bool) {
client, ok := h.provider.GetClient(accountID)
if !ok {
http.Error(w, "Client not found: "+string(accountID), http.StatusNotFound)
return
}
fullStatus, err := client.Status()
if err != nil {
http.Error(w, "Error getting status: "+err.Error(), http.StatusInternalServerError)
return
}
// Parse filter parameters
query := r.URL.Query()
statusFilter := query.Get("filter-by-status")
connectionTypeFilter := query.Get("filter-by-connection-type")
var prefixNamesFilter []string
var prefixNamesFilterMap map[string]struct{}
if names := query.Get("filter-by-names"); names != "" {
prefixNamesFilter = strings.Split(names, ",")
prefixNamesFilterMap = make(map[string]struct{})
for _, name := range prefixNamesFilter {
prefixNamesFilterMap[strings.ToLower(strings.TrimSpace(name))] = struct{}{}
}
}
var ipsFilterMap map[string]struct{}
if ips := query.Get("filter-by-ips"); ips != "" {
ipsFilterMap = make(map[string]struct{})
for _, ip := range strings.Split(ips, ",") {
ipsFilterMap[strings.TrimSpace(ip)] = struct{}{}
}
}
pbStatus := nbstatus.ToProtoFullStatus(fullStatus)
overview := nbstatus.ConvertToStatusOutputOverview(pbStatus, nbstatus.ConvertOptions{
StatusFilter: statusFilter,
PrefixNamesFilter: prefixNamesFilter,
PrefixNamesFilterMap: prefixNamesFilterMap,
IPsFilter: ipsFilterMap,
ConnectionTypeFilter: connectionTypeFilter,
})
if wantJSON {
resp := map[string]any{
"account_id": accountID,
"status": overview.FullDetailSummary(),
}
if info, ok := h.inboundInfoFor(accountID); ok {
resp["inbound_listener"] = info
}
h.writeJSON(w, resp)
return
}
data := clientDetailData{
AccountID: string(accountID),
ActiveTab: "status",
Content: overview.FullDetailSummary(),
}
h.renderTemplate(w, "clientDetail", data)
}
// inboundInfoFor returns the inbound listener info for an account, or
// ok=false when no inbound provider is wired or the account has no live
// listener.
func (h *Handler) inboundInfoFor(accountID types.AccountID) (InboundListenerInfo, bool) {
if h.inbound == nil {
return InboundListenerInfo{}, false
}
all := h.inbound.InboundListeners()
info, ok := all[accountID]
return info, ok
}
func (h *Handler) handleClientSyncResponse(w http.ResponseWriter, _ *http.Request, accountID types.AccountID, wantJSON bool) {
client, ok := h.provider.GetClient(accountID)
if !ok {
http.Error(w, "Client not found: "+string(accountID), http.StatusNotFound)
return
}
syncResp, err := client.GetLatestSyncResponse()
if err != nil {
http.Error(w, "Error getting sync response: "+err.Error(), http.StatusInternalServerError)
return
}
if syncResp == nil {
http.Error(w, "No sync response available for client: "+string(accountID), http.StatusNotFound)
return
}
opts := protojson.MarshalOptions{
EmitUnpopulated: true,
UseProtoNames: true,
Indent: " ",
AllowPartial: true,
}
jsonBytes, err := opts.Marshal(syncResp)
if err != nil {
http.Error(w, "Error marshaling sync response: "+err.Error(), http.StatusInternalServerError)
return
}
if wantJSON {
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write(jsonBytes)
return
}
data := clientDetailData{
AccountID: string(accountID),
ActiveTab: "syncresponse",
Content: string(jsonBytes),
}
h.renderTemplate(w, "clientDetail", data)
}
type toolsData struct {
AccountID string
}
func (h *Handler) handleClientTools(w http.ResponseWriter, _ *http.Request, accountID types.AccountID) {
_, ok := h.provider.GetClient(accountID)
if !ok {
http.Error(w, "Client not found: "+string(accountID), http.StatusNotFound)
return
}
data := toolsData{
AccountID: string(accountID),
}
h.renderTemplate(w, "tools", data)
}
func (h *Handler) handlePingTCP(w http.ResponseWriter, r *http.Request, accountID types.AccountID) {
client, ok := h.provider.GetClient(accountID)
if !ok {
h.writeJSON(w, map[string]any{"error": "client not found"})
return
}
host := r.URL.Query().Get("host")
portStr := r.URL.Query().Get("port")
if host == "" || portStr == "" {
h.writeJSON(w, map[string]any{"error": "host and port parameters required"})
return
}
port, err := strconv.Atoi(portStr)
if err != nil || port < 1 || port > 65535 {
h.writeJSON(w, map[string]any{"error": "invalid port"})
return
}
timeout := defaultPingTimeout
if t := r.URL.Query().Get("timeout"); t != "" {
if d, err := time.ParseDuration(t); err == nil {
timeout = d
}
}
network := "tcp"
if v := r.URL.Query().Get("ip_version"); v == "4" || v == "6" {
network += v
}
ctx, cancel := context.WithTimeout(r.Context(), timeout)
defer cancel()
address := net.JoinHostPort(host, strconv.Itoa(port))
start := time.Now()
conn, err := client.Dial(ctx, network, address)
if err != nil {
h.writeJSON(w, map[string]any{
"success": false,
"host": host,
"port": port,
"error": err.Error(),
})
return
}
remote := conn.RemoteAddr().String()
if err := conn.Close(); err != nil {
h.logger.Debugf("close tcp ping connection: %v", err)
}
latency := time.Since(start)
h.writeJSON(w, map[string]any{
"success": true,
"host": host,
"port": port,
"remote": remote,
"latency_ms": latency.Milliseconds(),
"latency": formatDuration(latency),
})
}
func (h *Handler) handleLogLevel(w http.ResponseWriter, r *http.Request, accountID types.AccountID) {
client, ok := h.provider.GetClient(accountID)
if !ok {
h.writeJSON(w, map[string]any{"error": "client not found"})
return
}
level := r.URL.Query().Get("level")
if level == "" {
h.writeJSON(w, map[string]any{"error": "level parameter required (trace, debug, info, warn, error)"})
return
}
if err := client.SetLogLevel(level); err != nil {
h.writeJSON(w, map[string]any{
"success": false,
"error": err.Error(),
})
return
}
h.writeJSON(w, map[string]any{
"success": true,
"level": level,
})
}
const clientActionTimeout = 30 * time.Second
func (h *Handler) handleClientStart(w http.ResponseWriter, r *http.Request, accountID types.AccountID) {
client, ok := h.provider.GetClient(accountID)
if !ok {
h.writeJSON(w, map[string]any{"error": "client not found"})
return
}
ctx, cancel := context.WithTimeout(r.Context(), clientActionTimeout)
defer cancel()
if err := client.Start(ctx); err != nil {
h.writeJSON(w, map[string]any{
"success": false,
"error": err.Error(),
})
return
}
h.writeJSON(w, map[string]any{
"success": true,
"message": "client started",
})
}
func (h *Handler) handleClientStop(w http.ResponseWriter, r *http.Request, accountID types.AccountID) {
client, ok := h.provider.GetClient(accountID)
if !ok {
h.writeJSON(w, map[string]any{"error": "client not found"})
return
}
ctx, cancel := context.WithTimeout(r.Context(), clientActionTimeout)
defer cancel()
if err := client.Stop(ctx); err != nil {
h.writeJSON(w, map[string]any{
"success": false,
"error": err.Error(),
})
return
}
h.writeJSON(w, map[string]any{
"success": true,
"message": "client stopped",
})
}
func (h *Handler) handlePerf(w http.ResponseWriter, r *http.Request) {
raw := r.URL.Query().Get("value")
if raw == "" {
http.Error(w, "value parameter is required", http.StatusBadRequest)
return
}
n, err := strconv.ParseUint(raw, 10, 32)
if err != nil {
http.Error(w, fmt.Sprintf("invalid value %q: %v", raw, err), http.StatusBadRequest)
return
}
capN := uint32(n)
applied := 0
failed := map[string]string{}
for accountID, client := range h.provider.ListClientsForStartup() {
if err := client.SetPerformance(nbembed.Performance{PreallocatedBuffersPerPool: &capN}); err != nil {
failed[string(accountID)] = err.Error()
continue
}
applied++
}
resp := map[string]any{
"success": true,
"value": capN,
"applied": applied,
}
if len(failed) > 0 {
resp["failed"] = failed
}
h.writeJSON(w, resp)
}
// handleRuntime returns cheap runtime and process stats. Safe to hit on a
// running proxy; does not read pprof profiles.
func (h *Handler) handleRuntime(w http.ResponseWriter, _ *http.Request) {
var m runtime.MemStats
runtime.ReadMemStats(&m)
clients := h.provider.ListClientsForDebug()
started := 0
for _, c := range clients {
if c.HasClient {
started++
}
}
resp := map[string]any{
"uptime": time.Since(h.startTime).Round(time.Second).String(),
"goroutines": runtime.NumGoroutine(),
"num_cpu": runtime.NumCPU(),
"gomaxprocs": runtime.GOMAXPROCS(0),
"go_version": runtime.Version(),
"heap_alloc": m.HeapAlloc,
"heap_inuse": m.HeapInuse,
"heap_idle": m.HeapIdle,
"heap_released": m.HeapReleased,
"heap_sys": m.HeapSys,
"sys": m.Sys,
"live_objects": m.Mallocs - m.Frees,
"num_gc": m.NumGC,
"pause_total_ns": m.PauseTotalNs,
"clients": len(clients),
"started": started,
}
if proc := readProcStatus(); proc != nil {
resp["vm_rss"] = proc["VmRSS"]
resp["vm_size"] = proc["VmSize"]
resp["vm_data"] = proc["VmData"]
}
h.writeJSON(w, resp)
}
// readProcStatus parses /proc/self/status on Linux and returns size fields
// in bytes. Returns nil on non-Linux or read failure.
func readProcStatus() map[string]uint64 {
raw, err := os.ReadFile("/proc/self/status")
if err != nil {
return nil
}
out := map[string]uint64{}
for _, line := range strings.Split(string(raw), "\n") {
k, v, ok := strings.Cut(line, ":")
if !ok {
continue
}
if k != "VmRSS" && k != "VmSize" && k != "VmData" {
continue
}
fields := strings.Fields(v)
if len(fields) < 1 {
continue
}
n, err := strconv.ParseUint(fields[0], 10, 64)
if err != nil {
continue
}
// Values are reported in kB.
out[k] = n * 1024
}
return out
}
const maxCaptureDuration = 30 * time.Minute
// handleCapture streams a pcap or text packet capture for the given client.
//
// Query params:
//
// duration: capture duration (0 or absent = max, capped at 30m)
// format: "text" for human-readable output (default: pcap)
// filter: BPF-like filter expression (e.g. "host 10.0.0.1 and tcp port 443")
func (h *Handler) handleCapture(w http.ResponseWriter, r *http.Request, accountID types.AccountID) {
client, ok := h.provider.GetClient(accountID)
if !ok {
http.Error(w, "client not found", http.StatusNotFound)
return
}
duration := maxCaptureDuration
if durationStr := r.URL.Query().Get("duration"); durationStr != "" {
d, err := time.ParseDuration(durationStr)
if err != nil {
http.Error(w, "invalid duration: "+err.Error(), http.StatusBadRequest)
return
}
if d < 0 {
http.Error(w, "duration must not be negative", http.StatusBadRequest)
return
}
if d > 0 {
duration = min(d, maxCaptureDuration)
}
}
filter := r.URL.Query().Get("filter")
wantText := r.URL.Query().Get("format") == "text"
verbose := r.URL.Query().Get("verbose") == "true"
ascii := r.URL.Query().Get("ascii") == "true"
opts := nbembed.CaptureOptions{Filter: filter, Verbose: verbose, ASCII: ascii}
if wantText {
w.Header().Set("Content-Type", "text/plain; charset=utf-8")
opts.TextOutput = w
} else {
w.Header().Set("Content-Type", "application/vnd.tcpdump.pcap")
w.Header().Set("Content-Disposition",
fmt.Sprintf("attachment; filename=capture-%s.pcap", accountID))
opts.Output = w
}
cs, err := client.StartCapture(opts)
if err != nil {
http.Error(w, "start capture: "+err.Error(), http.StatusServiceUnavailable)
return
}
defer cs.Stop()
// Flush headers after setup succeeds so errors above can still set status codes.
if f, ok := w.(http.Flusher); ok {
f.Flush()
}
timer := time.NewTimer(duration)
defer timer.Stop()
select {
case <-r.Context().Done():
case <-timer.C:
}
cs.Stop()
stats := cs.Stats()
h.logger.Infof("capture for %s finished: %d packets, %d bytes, %d dropped",
accountID, stats.Packets, stats.Bytes, stats.Dropped)
}
func (h *Handler) handleHealth(w http.ResponseWriter, r *http.Request, wantJSON bool) {
if !wantJSON {
http.Redirect(w, r, "/debug", http.StatusSeeOther)
return
}
uptime := time.Since(h.startTime).Round(10 * time.Millisecond).String()
ready := h.health.ReadinessProbe()
allHealthy, clientHealth := h.health.CheckClientsConnected(r.Context())
status := "ok"
// No clients is not a health issue; only degrade when actual clients are unhealthy
if !ready || (!allHealthy && len(clientHealth) > 0) {
status = "degraded"
}
var certsTotal, certsReady, certsPending, certsFailed int
var certsPendingDomains, certsReadyDomains []string
var certsFailedDomains map[string]string
if h.certStatus != nil {
certsTotal = h.certStatus.TotalDomains()
certsPendingDomains = h.certStatus.PendingDomains()
certsReadyDomains = h.certStatus.ReadyDomains()
certsFailedDomains = h.certStatus.FailedDomains()
certsReady = len(certsReadyDomains)
certsPending = len(certsPendingDomains)
certsFailed = len(certsFailedDomains)
}
resp := map[string]any{
"status": status,
"uptime": uptime,
"management_connected": ready,
"all_clients_healthy": allHealthy,
"certs_total": certsTotal,
"certs_ready": certsReady,
"certs_pending": certsPending,
"certs_failed": certsFailed,
"clients": clientHealth,
}
if len(certsPendingDomains) > 0 {
resp["certs_pending_domains"] = certsPendingDomains
}
if len(certsReadyDomains) > 0 {
resp["certs_ready_domains"] = certsReadyDomains
}
if len(certsFailedDomains) > 0 {
resp["certs_failed_domains"] = certsFailedDomains
}
h.writeJSON(w, resp)
}
func (h *Handler) renderTemplate(w http.ResponseWriter, name string, data any) {
w.Header().Set("Content-Type", "text/html; charset=utf-8")
tmpl := h.getTemplates()
if tmpl == nil {
http.Error(w, "Templates not loaded", http.StatusInternalServerError)
return
}
if err := tmpl.ExecuteTemplate(w, name, data); err != nil {
h.logger.Errorf("execute template %s: %v", name, err)
http.Error(w, "Template error", http.StatusInternalServerError)
}
}
func (h *Handler) writeJSON(w http.ResponseWriter, v any) {
w.Header().Set("Content-Type", "application/json")
enc := json.NewEncoder(w)
enc.SetIndent("", " ")
if err := enc.Encode(v); err != nil {
h.logger.Errorf("encode JSON response: %v", err)
}
}
proxy/internal/debug/templates/base.html
+101
View File
@@ -0,0 +1,101 @@
{{define "style"}}
body {
font-family: monospace;
margin: 20px;
background: #1a1a1a;
color: #eee;
}
a {
color: #6cf;
}
h1, h2, h3 {
color: #fff;
}
.info {
color: #aaa;
}
table {
border-collapse: collapse;
margin: 10px 0;
}
th, td {
border: 1px solid #444;
padding: 8px;
text-align: left;
}
th {
background: #333;
}
.nav {
margin-bottom: 20px;
}
.nav a {
margin-right: 15px;
padding: 8px 16px;
background: #333;
text-decoration: none;
border-radius: 4px;
}
.nav a.active {
background: #6cf;
color: #000;
}
pre {
background: #222;
padding: 15px;
border-radius: 4px;
overflow-x: auto;
white-space: pre-wrap;
}
input, select, textarea {
background: #333;
color: #eee;
border: 1px solid #555;
padding: 8px;
border-radius: 4px;
font-family: monospace;
}
input:focus, select:focus, textarea:focus {
outline: none;
border-color: #6cf;
}
button {
background: #6cf;
color: #000;
border: none;
padding: 8px 16px;
border-radius: 4px;
cursor: pointer;
font-family: monospace;
}
button:hover {
background: #5be;
}
button:disabled {
background: #555;
color: #888;
cursor: not-allowed;
}
.form-group {
margin-bottom: 15px;
}
.form-group label {
display: block;
margin-bottom: 5px;
color: #aaa;
}
.form-row {
display: flex;
gap: 10px;
align-items: flex-end;
}
.result {
margin-top: 20px;
}
.success {
color: #5f5;
}
.error {
color: #f55;
}
{{end}}
proxy/internal/debug/templates/client_detail.html
+19
View File
@@ -0,0 +1,19 @@
{{define "clientDetail"}}
<!DOCTYPE html>
<html lang="en">
<head>
<title>Client {{.AccountID}}</title>
<style>{{template "style"}}</style>
</head>
<body>
<h1>Client: {{.AccountID}}</h1>
<div class="nav">
<a href="/debug">&larr; Back</a>
<a href="/debug/clients/{{.AccountID}}/tools"{{if eq .ActiveTab "tools"}} class="active"{{end}}>Tools</a>
<a href="/debug/clients/{{.AccountID}}"{{if eq .ActiveTab "status"}} class="active"{{end}}>Status</a>
<a href="/debug/clients/{{.AccountID}}/syncresponse"{{if eq .ActiveTab "syncresponse"}} class="active"{{end}}>Sync Response</a>
</div>
<pre>{{.Content}}</pre>
</body>
</html>
{{end}}
proxy/internal/debug/templates/clients.html
+33
View File
@@ -0,0 +1,33 @@
{{define "clients"}}
<!DOCTYPE html>
<html lang="en">
<head>
<title>Clients</title>
<style>{{template "style"}}</style>
</head>
<body>
<h1>All Clients</h1>
<p class="info">Uptime: {{.Uptime}} | <a href="/debug">&larr; Back</a></p>
{{if .Clients}}
<table>
<tr>
<th>Account ID</th>
<th>Services</th>
<th>Age</th>
<th>Status</th>
</tr>
{{range .Clients}}
<tr>
<td><a href="/debug/clients/{{.AccountID}}/tools">{{.AccountID}}</a></td>
<td>{{.Services}}</td>
<td>{{.Age}}</td>
<td>{{.Status}}</td>
</tr>
{{end}}
</table>
{{else}}
<p>No clients connected</p>
{{end}}
</body>
</html>
{{end}}
proxy/internal/debug/templates/index.html
+58
View File
@@ -0,0 +1,58 @@
{{define "index"}}
<!DOCTYPE html>
<html lang="en">
<head>
<title>NetBird Proxy Debug</title>
<style>{{template "style"}}</style>
</head>
<body>
<h1>NetBird Proxy Debug</h1>
<p class="info">Version: {{.Version}} | Uptime: {{.Uptime}}</p>
<h2>Certificates: {{.CertsReady}} ready, {{.CertsPending}} pending, {{.CertsFailed}} failed ({{.CertsTotal}} total)</h2>
{{if .CertsReadyDomains}}
<details>
<summary>Ready domains ({{.CertsReady}})</summary>
<ul>{{range .CertsReadyDomains}}<li>{{.}}</li>{{end}}</ul>
</details>
{{end}}
{{if .CertsPendingDomains}}
<details open>
<summary>Pending domains ({{.CertsPending}})</summary>
<ul>{{range .CertsPendingDomains}}<li>{{.}}</li>{{end}}</ul>
</details>
{{end}}
{{if .CertsFailedDomains}}
<details open>
<summary>Failed domains ({{.CertsFailed}})</summary>
<ul>{{range .CertsFailedDomains}}<li>{{.Domain}}: {{.Error}}</li>{{end}}</ul>
</details>
{{end}}
<h2>Clients ({{.ClientCount}}) | Services ({{.TotalServices}})</h2>
{{if .Clients}}
<table>
<tr>
<th>Account ID</th>
<th>Services</th>
<th>Age</th>
<th>Status</th>
</tr>
{{range .Clients}}
<tr>
<td><a href="/debug/clients/{{.AccountID}}/tools">{{.AccountID}}</a></td>
<td>{{.Services}}</td>
<td>{{.Age}}</td>
<td>{{.Status}}</td>
</tr>
{{end}}
</table>
{{else}}
<p>No clients connected</p>
{{end}}
<h2>Endpoints</h2>
<ul>
<li><a href="/debug/clients">/debug/clients</a> - all clients detail</li>
</ul>
<p class="info">Add ?format=json or /json suffix for JSON output</p>
</body>
</html>
{{end}}
proxy/internal/debug/templates/tools.html
+142
View File
@@ -0,0 +1,142 @@
{{define "tools"}}
<!DOCTYPE html>
<html lang="en">
<head>
<title>Client {{.AccountID}} - Tools</title>
<style>{{template "style"}}</style>
</head>
<body>
<h1>Client: {{.AccountID}}</h1>
<div class="nav">
<a href="/debug">&larr; Back</a>
<a href="/debug/clients/{{.AccountID}}/tools" class="active">Tools</a>
<a href="/debug/clients/{{.AccountID}}">Status</a>
<a href="/debug/clients/{{.AccountID}}/syncresponse">Sync Response</a>
</div>
<h2>Client Control</h2>
<div class="form-row">
<div class="form-group">
<span>&nbsp;</span>
<button onclick="startClient()">Start</button>
</div>
<div class="form-group">
<span>&nbsp;</span>
<button onclick="stopClient()">Stop</button>
</div>
</div>
<div id="client-result" class="result"></div>
<h2>Log Level</h2>
<div class="form-row">
<div class="form-group">
<label for="log-level">Level</label>
<select id="log-level" style="width: 120px;">
<option value="trace">trace</option>
<option value="debug">debug</option>
<option value="info">info</option>
<option value="warn" selected>warn</option>
<option value="error">error</option>
</select>
</div>
<div class="form-group">
<span>&nbsp;</span>
<button onclick="setLogLevel()">Set Level</button>
</div>
</div>
<div id="log-result" class="result"></div>
<h2>TCP Ping</h2>
<div class="form-row">
<div class="form-group">
<label for="tcp-host">Host</label>
<input type="text" id="tcp-host" placeholder="100.0.0.1 or hostname.netbird.cloud" style="width: 300px;">
</div>
<div class="form-group">
<label for="tcp-port">Port</label>
<input type="number" id="tcp-port" placeholder="80" style="width: 80px;">
</div>
<div class="form-group">
<span>&nbsp;</span>
<button onclick="doTcpPing()">Connect</button>
</div>
</div>
<div id="tcp-result" class="result"></div>
<script>
const accountID = "{{.AccountID}}";
async function startClient() {
const resultDiv = document.getElementById('client-result');
resultDiv.innerHTML = '<span class="info">Starting client...</span>';
try {
const resp = await fetch('/debug/clients/' + accountID + '/start');
const data = await resp.json();
if (data.success) {
resultDiv.innerHTML = '<span class="success">✓ ' + data.message + '</span>';
} else {
resultDiv.innerHTML = '<span class="error">✗ ' + data.error + '</span>';
}
} catch (e) {
resultDiv.innerHTML = '<span class="error">Error: ' + e.message + '</span>';
}
}
async function stopClient() {
const resultDiv = document.getElementById('client-result');
resultDiv.innerHTML = '<span class="info">Stopping client...</span>';
try {
const resp = await fetch('/debug/clients/' + accountID + '/stop');
const data = await resp.json();
if (data.success) {
resultDiv.innerHTML = '<span class="success">✓ ' + data.message + '</span>';
} else {
resultDiv.innerHTML = '<span class="error">✗ ' + data.error + '</span>';
}
} catch (e) {
resultDiv.innerHTML = '<span class="error">Error: ' + e.message + '</span>';
}
}
async function setLogLevel() {
const level = document.getElementById('log-level').value;
const resultDiv = document.getElementById('log-result');
resultDiv.innerHTML = '<span class="info">Setting log level...</span>';
try {
const resp = await fetch('/debug/clients/' + accountID + '/loglevel?level=' + level);
const data = await resp.json();
if (data.success) {
resultDiv.innerHTML = '<span class="success">✓ Log level set to: ' + data.level + '</span>';
} else {
resultDiv.innerHTML = '<span class="error">✗ ' + data.error + '</span>';
}
} catch (e) {
resultDiv.innerHTML = '<span class="error">Error: ' + e.message + '</span>';
}
}
async function doTcpPing() {
const host = document.getElementById('tcp-host').value;
const port = document.getElementById('tcp-port').value;
if (!host || !port) {
alert('Host and port required');
return;
}
const resultDiv = document.getElementById('tcp-result');
resultDiv.innerHTML = '<span class="info">Connecting...</span>';
try {
const resp = await fetch('/debug/clients/' + accountID + '/pingtcp?host=' + encodeURIComponent(host) + '&port=' + port);
const data = await resp.json();
if (data.success) {
resultDiv.innerHTML = '<span class="success">✓ ' + data.host + ':' + data.port + ' connected in ' + data.latency + '</span>';
} else {
resultDiv.innerHTML = '<span class="error">✗ ' + data.host + ':' + data.port + ': ' + data.error + '</span>';
}
} catch (e) {
resultDiv.innerHTML = '<span class="error">Error: ' + e.message + '</span>';
}
}
</script>
</body>
</html>
{{end}}
proxy/internal/flock/flock_other.go
+20
View File
@@ -0,0 +1,20 @@
//go:build !unix
package flock
import (
"context"
"os"
)
// Lock is a no-op on non-Unix platforms. Returns (nil, nil) to indicate
// that no lock was acquired; callers must treat a nil file as "proceed
// without lock" rather than "lock held by someone else."
func Lock(_ context.Context, _ string) (*os.File, error) {
return nil, nil //nolint:nilnil // intentional: nil file signals locking unsupported on this platform
}
// Unlock is a no-op on non-Unix platforms.
func Unlock(_ *os.File) error {
return nil
}
proxy/internal/flock/flock_test.go
+79
View File
@@ -0,0 +1,79 @@
//go:build unix
package flock
import (
"context"
"os"
"path/filepath"
"sync"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestLockUnlock(t *testing.T) {
lockPath := filepath.Join(t.TempDir(), "test.lock")
f, err := Lock(context.Background(), lockPath)
require.NoError(t, err)
require.NotNil(t, f)
_, err = os.Stat(lockPath)
assert.NoError(t, err, "lock file should exist")
err = Unlock(f)
assert.NoError(t, err)
}
func TestUnlockNil(t *testing.T) {
err := Unlock(nil)
assert.NoError(t, err, "unlocking nil should be a no-op")
}
func TestLockRespectsContext(t *testing.T) {
lockPath := filepath.Join(t.TempDir(), "test.lock")
f1, err := Lock(context.Background(), lockPath)
require.NoError(t, err)
defer func() { require.NoError(t, Unlock(f1)) }()
ctx, cancel := context.WithTimeout(context.Background(), 300*time.Millisecond)
defer cancel()
_, err = Lock(ctx, lockPath)
assert.ErrorIs(t, err, context.DeadlineExceeded)
}
func TestLockBlocks(t *testing.T) {
lockPath := filepath.Join(t.TempDir(), "test.lock")
f1, err := Lock(context.Background(), lockPath)
require.NoError(t, err)
var wg sync.WaitGroup
wg.Add(1)
start := time.Now()
var elapsed time.Duration
go func() {
defer wg.Done()
f2, err := Lock(context.Background(), lockPath)
elapsed = time.Since(start)
assert.NoError(t, err)
if f2 != nil {
assert.NoError(t, Unlock(f2))
}
}()
// Hold the lock for 200ms, then release.
time.Sleep(200 * time.Millisecond)
require.NoError(t, Unlock(f1))
wg.Wait()
assert.GreaterOrEqual(t, elapsed, 150*time.Millisecond,
"Lock should have blocked for at least ~200ms")
}
proxy/internal/flock/flock_unix.go
+77
View File
@@ -0,0 +1,77 @@
//go:build unix
// Package flock provides best-effort advisory file locking using flock(2).
//
// This is used for cross-replica coordination (e.g. preventing duplicate
// ACME requests). Note that flock(2) does NOT work reliably on NFS volumes:
// on NFSv3 it depends on the NLM daemon, on NFSv4 Linux emulates it via
// fcntl locks with different semantics. Callers must treat lock failures
// as non-fatal and proceed without the lock.
package flock
import (
"context"
"errors"
"fmt"
"os"
"syscall"
"time"
log "github.com/sirupsen/logrus"
)
const retryInterval = 100 * time.Millisecond
// Lock acquires an exclusive advisory lock on the given file path.
// It creates the lock file if it does not exist. The lock attempt
// respects context cancellation by using non-blocking flock with polling.
// The caller must call Unlock with the returned *os.File when done.
func Lock(ctx context.Context, path string) (*os.File, error) {
f, err := os.OpenFile(path, os.O_CREATE|os.O_RDWR, 0o600)
if err != nil {
return nil, fmt.Errorf("open lock file %s: %w", path, err)
}
timer := time.NewTimer(retryInterval)
defer timer.Stop()
for {
if err := syscall.Flock(int(f.Fd()), syscall.LOCK_EX|syscall.LOCK_NB); err == nil {
return f, nil
} else if !errors.Is(err, syscall.EWOULDBLOCK) {
if cerr := f.Close(); cerr != nil {
log.Debugf("close lock file %s: %v", path, cerr)
}
return nil, fmt.Errorf("acquire lock on %s: %w", path, err)
}
select {
case <-ctx.Done():
if cerr := f.Close(); cerr != nil {
log.Debugf("close lock file %s: %v", path, cerr)
}
return nil, ctx.Err()
case <-timer.C:
timer.Reset(retryInterval)
}
}
}
// Unlock releases the lock and closes the file.
func Unlock(f *os.File) error {
if f == nil {
return nil
}
defer func() {
if cerr := f.Close(); cerr != nil {
log.Debugf("close lock file: %v", cerr)
}
}()
if err := syscall.Flock(int(f.Fd()), syscall.LOCK_UN); err != nil {
return fmt.Errorf("release lock: %w", err)
}
return nil
}
proxy/internal/geolocation/download.go
+264
View File
@@ -0,0 +1,264 @@
package geolocation
import (
"archive/tar"
"bufio"
"compress/gzip"
"crypto/sha256"
"errors"
"fmt"
"io"
"mime"
"net/http"
"os"
"path/filepath"
"strings"
"time"
log "github.com/sirupsen/logrus"
)
const (
mmdbTarGZURL = "https://pkgs.netbird.io/geolocation-dbs/GeoLite2-City/download?suffix=tar.gz"
mmdbSha256URL = "https://pkgs.netbird.io/geolocation-dbs/GeoLite2-City/download?suffix=tar.gz.sha256"
mmdbInnerName = "GeoLite2-City.mmdb"
downloadTimeout = 2 * time.Minute
maxMMDBSize = 256 << 20 // 256 MB
)
// ensureMMDB checks for an existing MMDB file in dataDir. If none is found,
// it downloads from pkgs.netbird.io with SHA256 verification.
func ensureMMDB(logger *log.Logger, dataDir string) (string, error) {
if err := os.MkdirAll(dataDir, 0o755); err != nil {
return "", fmt.Errorf("create geo data directory %s: %w", dataDir, err)
}
pattern := filepath.Join(dataDir, mmdbGlob)
if files, _ := filepath.Glob(pattern); len(files) > 0 {
mmdbPath := files[len(files)-1]
logger.Debugf("using existing geolocation database: %s", mmdbPath)
return mmdbPath, nil
}
logger.Info("geolocation database not found, downloading from pkgs.netbird.io")
return downloadMMDB(logger, dataDir)
}
func downloadMMDB(logger *log.Logger, dataDir string) (string, error) {
client := &http.Client{Timeout: downloadTimeout}
datedName, err := fetchRemoteFilename(client, mmdbTarGZURL)
if err != nil {
return "", fmt.Errorf("get remote filename: %w", err)
}
mmdbFilename := deriveMMDBFilename(datedName)
mmdbPath := filepath.Join(dataDir, mmdbFilename)
tmp, err := os.MkdirTemp("", "geolite-proxy-*")
if err != nil {
return "", fmt.Errorf("create temp directory: %w", err)
}
defer os.RemoveAll(tmp)
checksumFile := filepath.Join(tmp, "checksum.sha256")
if err := downloadToFile(client, mmdbSha256URL, checksumFile); err != nil {
return "", fmt.Errorf("download checksum: %w", err)
}
expectedHash, err := readChecksumFile(checksumFile)
if err != nil {
return "", fmt.Errorf("read checksum: %w", err)
}
tarFile := filepath.Join(tmp, datedName)
logger.Debugf("downloading geolocation database (%s)", datedName)
if err := downloadToFile(client, mmdbTarGZURL, tarFile); err != nil {
return "", fmt.Errorf("download database: %w", err)
}
if err := verifySHA256(tarFile, expectedHash); err != nil {
return "", fmt.Errorf("verify database checksum: %w", err)
}
if err := extractMMDBFromTarGZ(tarFile, mmdbPath); err != nil {
return "", fmt.Errorf("extract database: %w", err)
}
logger.Infof("geolocation database downloaded: %s", mmdbPath)
return mmdbPath, nil
}
// deriveMMDBFilename converts a tar.gz filename to an MMDB filename.
// Example: GeoLite2-City_20240101.tar.gz -> GeoLite2-City_20240101.mmdb
func deriveMMDBFilename(tarName string) string {
base, _, _ := strings.Cut(tarName, ".")
if !strings.Contains(base, "_") {
return "GeoLite2-City.mmdb"
}
return base + ".mmdb"
}
func fetchRemoteFilename(client *http.Client, url string) (string, error) {
resp, err := client.Head(url)
if err != nil {
return "", err
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
return "", fmt.Errorf("HEAD request: HTTP %d", resp.StatusCode)
}
cd := resp.Header.Get("Content-Disposition")
if cd == "" {
return "", errors.New("no Content-Disposition header")
}
_, params, err := mime.ParseMediaType(cd)
if err != nil {
return "", fmt.Errorf("parse Content-Disposition: %w", err)
}
name := filepath.Base(params["filename"])
if name == "" || name == "." {
return "", errors.New("no filename in Content-Disposition")
}
return name, nil
}
func downloadToFile(client *http.Client, url, dest string) error {
resp, err := client.Get(url) //nolint:gosec
if err != nil {
return err
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
body, _ := io.ReadAll(io.LimitReader(resp.Body, 1024))
return fmt.Errorf("HTTP %d: %s", resp.StatusCode, string(body))
}
f, err := os.Create(dest) //nolint:gosec
if err != nil {
return err
}
defer f.Close()
// Cap download at 256 MB to prevent unbounded reads from a compromised server.
if _, err := io.Copy(f, io.LimitReader(resp.Body, maxMMDBSize)); err != nil {
return err
}
return nil
}
func readChecksumFile(path string) (string, error) {
f, err := os.Open(path) //nolint:gosec
if err != nil {
return "", err
}
defer f.Close()
scanner := bufio.NewScanner(f)
if scanner.Scan() {
parts := strings.Fields(scanner.Text())
if len(parts) > 0 {
return parts[0], nil
}
}
if err := scanner.Err(); err != nil {
return "", err
}
return "", errors.New("empty checksum file")
}
func verifySHA256(path, expected string) error {
f, err := os.Open(path) //nolint:gosec
if err != nil {
return err
}
defer f.Close()
h := sha256.New()
if _, err := io.Copy(h, f); err != nil {
return err
}
actual := fmt.Sprintf("%x", h.Sum(nil))
if actual != expected {
return fmt.Errorf("SHA256 mismatch: expected %s, got %s", expected, actual)
}
return nil
}
func extractMMDBFromTarGZ(tarGZPath, destPath string) error {
f, err := os.Open(tarGZPath) //nolint:gosec
if err != nil {
return err
}
defer f.Close()
gz, err := gzip.NewReader(f)
if err != nil {
return err
}
defer gz.Close()
tr := tar.NewReader(gz)
for {
hdr, err := tr.Next()
if err != nil {
if errors.Is(err, io.EOF) {
break
}
return err
}
if hdr.Typeflag == tar.TypeReg && filepath.Base(hdr.Name) == mmdbInnerName {
if hdr.Size < 0 || hdr.Size > maxMMDBSize {
return fmt.Errorf("mmdb entry size %d exceeds limit %d", hdr.Size, maxMMDBSize)
}
if err := extractToFileAtomic(io.LimitReader(tr, hdr.Size), destPath); err != nil {
return err
}
return nil
}
}
return fmt.Errorf("%s not found in archive", mmdbInnerName)
}
// extractToFileAtomic writes r to a temporary file in the same directory as
// destPath, then renames it into place so a crash never leaves a truncated file.
func extractToFileAtomic(r io.Reader, destPath string) error {
dir := filepath.Dir(destPath)
tmp, err := os.CreateTemp(dir, ".mmdb-*.tmp")
if err != nil {
return fmt.Errorf("create temp file: %w", err)
}
tmpPath := tmp.Name()
if _, err := io.Copy(tmp, r); err != nil { //nolint:gosec // G110: caller bounds with LimitReader
if closeErr := tmp.Close(); closeErr != nil {
log.Debugf("failed to close temp file %s: %v", tmpPath, closeErr)
}
if removeErr := os.Remove(tmpPath); removeErr != nil {
log.Debugf("failed to remove temp file %s: %v", tmpPath, removeErr)
}
return fmt.Errorf("write mmdb: %w", err)
}
if err := tmp.Close(); err != nil {
if removeErr := os.Remove(tmpPath); removeErr != nil {
log.Debugf("failed to remove temp file %s: %v", tmpPath, removeErr)
}
return fmt.Errorf("close temp file: %w", err)
}
if err := os.Rename(tmpPath, destPath); err != nil {
if removeErr := os.Remove(tmpPath); removeErr != nil {
log.Debugf("failed to remove temp file %s: %v", tmpPath, removeErr)
}
return fmt.Errorf("rename to %s: %w", destPath, err)
}
return nil
}
proxy/internal/geolocation/geolocation.go
+152
View File
@@ -0,0 +1,152 @@
// Package geolocation provides IP-to-country lookups using MaxMind GeoLite2 databases.
package geolocation
import (
"fmt"
"net/netip"
"os"
"strconv"
"sync"
"github.com/oschwald/maxminddb-golang"
log "github.com/sirupsen/logrus"
)
const (
// EnvDisable disables geolocation lookups entirely when set to a truthy value.
EnvDisable = "NB_PROXY_DISABLE_GEOLOCATION"
mmdbGlob = "GeoLite2-City_*.mmdb"
)
type record struct {
Country struct {
ISOCode string `maxminddb:"iso_code"`
} `maxminddb:"country"`
City struct {
Names struct {
En string `maxminddb:"en"`
} `maxminddb:"names"`
} `maxminddb:"city"`
Subdivisions []struct {
ISOCode string `maxminddb:"iso_code"`
Names struct {
En string `maxminddb:"en"`
} `maxminddb:"names"`
} `maxminddb:"subdivisions"`
}
// Result holds the outcome of a geo lookup.
type Result struct {
CountryCode string
CityName string
SubdivisionCode string
SubdivisionName string
}
// Lookup provides IP geolocation lookups.
type Lookup struct {
mu sync.RWMutex
db *maxminddb.Reader
logger *log.Logger
}
// NewLookup opens or downloads the GeoLite2-City MMDB in dataDir.
// Returns nil without error if geolocation is disabled via environment
// variable, no data directory is configured, or the download fails
// (graceful degradation: country restrictions will deny all requests).
func NewLookup(logger *log.Logger, dataDir string) (*Lookup, error) {
if isDisabledByEnv(logger) {
logger.Info("geolocation disabled via environment variable")
return nil, nil //nolint:nilnil
}
if dataDir == "" {
return nil, nil //nolint:nilnil
}
mmdbPath, err := ensureMMDB(logger, dataDir)
if err != nil {
logger.Warnf("geolocation database unavailable: %v", err)
logger.Warn("country-based access restrictions will deny all requests until a database is available")
return nil, nil //nolint:nilnil
}
db, err := maxminddb.Open(mmdbPath)
if err != nil {
return nil, fmt.Errorf("open GeoLite2 database %s: %w", mmdbPath, err)
}
logger.Infof("geolocation database loaded from %s", mmdbPath)
return &Lookup{db: db, logger: logger}, nil
}
// LookupAddr returns the country ISO code and city name for the given IP.
// Returns an empty Result if the database is nil or the lookup fails.
func (l *Lookup) LookupAddr(addr netip.Addr) Result {
if l == nil {
return Result{}
}
l.mu.RLock()
defer l.mu.RUnlock()
if l.db == nil {
return Result{}
}
addr = addr.Unmap()
var rec record
if err := l.db.Lookup(addr.AsSlice(), &rec); err != nil {
l.logger.Debugf("geolocation lookup %s: %v", addr, err)
return Result{}
}
r := Result{
CountryCode: rec.Country.ISOCode,
CityName: rec.City.Names.En,
}
if len(rec.Subdivisions) > 0 {
r.SubdivisionCode = rec.Subdivisions[0].ISOCode
r.SubdivisionName = rec.Subdivisions[0].Names.En
}
return r
}
// Available reports whether the lookup has a loaded database.
func (l *Lookup) Available() bool {
if l == nil {
return false
}
l.mu.RLock()
defer l.mu.RUnlock()
return l.db != nil
}
// Close releases the database resources.
func (l *Lookup) Close() error {
if l == nil {
return nil
}
l.mu.Lock()
defer l.mu.Unlock()
if l.db != nil {
err := l.db.Close()
l.db = nil
return err
}
return nil
}
func isDisabledByEnv(logger *log.Logger) bool {
val := os.Getenv(EnvDisable)
if val == "" {
return false
}
disabled, err := strconv.ParseBool(val)
if err != nil {
logger.Warnf("parse %s=%q: %v", EnvDisable, val, err)
return false
}
return disabled
}
proxy/internal/grpc/auth.go
+48
View File
@@ -0,0 +1,48 @@
// Package grpc provides gRPC utilities for the proxy client.
package grpc
import (
"context"
"os"
"strconv"
log "github.com/sirupsen/logrus"
"google.golang.org/grpc"
"google.golang.org/grpc/credentials"
)
// EnvProxyAllowInsecure controls whether the proxy token can be sent over non-TLS connections.
const EnvProxyAllowInsecure = "NB_PROXY_ALLOW_INSECURE"
var _ credentials.PerRPCCredentials = (*proxyAuthToken)(nil)
type proxyAuthToken struct {
token string
allowInsecure bool
}
func (t proxyAuthToken) GetRequestMetadata(context.Context, ...string) (map[string]string, error) {
return map[string]string{
"authorization": "Bearer " + t.token,
}, nil
}
// RequireTransportSecurity returns true by default to protect the token in transit.
// Set NB_PROXY_ALLOW_INSECURE=true to allow non-TLS connections (not recommended for production).
func (t proxyAuthToken) RequireTransportSecurity() bool {
return !t.allowInsecure
}
// WithProxyToken returns a DialOption that sets the proxy access token on each outbound RPC.
func WithProxyToken(token string) grpc.DialOption {
allowInsecure := false
if val := os.Getenv(EnvProxyAllowInsecure); val != "" {
parsed, err := strconv.ParseBool(val)
if err != nil {
log.Warnf("invalid value for %s: %v", EnvProxyAllowInsecure, err)
} else {
allowInsecure = parsed
}
}
return grpc.WithPerRPCCredentials(proxyAuthToken{token: token, allowInsecure: allowInsecure})
}
proxy/internal/health/health.go
+405
View File
@@ -0,0 +1,405 @@
// Package health provides health probes for the proxy server.
package health
import (
"context"
"encoding/json"
"net"
"net/http"
"strings"
"sync"
"time"
log "github.com/sirupsen/logrus"
"github.com/netbirdio/netbird/client/embed"
"github.com/netbirdio/netbird/proxy/internal/types"
)
const handshakeStaleThreshold = 5 * time.Minute
const (
maxConcurrentChecks = 3
maxClientCheckTimeout = 5 * time.Minute
)
// clientProvider provides access to NetBird clients for health checks.
type clientProvider interface {
ListClientsForStartup() map[types.AccountID]*embed.Client
}
// Checker tracks health state and provides probe endpoints.
type Checker struct {
logger *log.Logger
provider clientProvider
mu sync.RWMutex
managementConnected bool
initialSyncComplete bool
shuttingDown bool
// checkSem limits concurrent client health checks.
checkSem chan struct{}
// checkHealth checks the health of a single client.
// Defaults to checkClientHealth; overridable in tests.
checkHealth func(*embed.Client) ClientHealth
}
// ClientHealth represents the health status of a single NetBird client.
type ClientHealth struct {
Healthy bool `json:"healthy"`
ManagementConnected bool `json:"management_connected"`
SignalConnected bool `json:"signal_connected"`
RelaysConnected int `json:"relays_connected"`
RelaysTotal int `json:"relays_total"`
PeersTotal int `json:"peers_total"`
PeersConnected int `json:"peers_connected"`
PeersP2P int `json:"peers_p2p"`
PeersRelayed int `json:"peers_relayed"`
PeersDegraded int `json:"peers_degraded"`
Error string `json:"error,omitempty"`
}
// ProbeResponse represents the JSON response for health probes.
type ProbeResponse struct {
Status string `json:"status"`
Checks map[string]bool `json:"checks,omitempty"`
Clients map[types.AccountID]ClientHealth `json:"clients,omitempty"`
}
// Server runs the health probe HTTP server on a dedicated port.
type Server struct {
server *http.Server
logger *log.Logger
checker *Checker
}
// SetManagementConnected updates the management connection state.
func (c *Checker) SetManagementConnected(connected bool) {
c.mu.Lock()
defer c.mu.Unlock()
c.managementConnected = connected
}
// SetInitialSyncComplete marks that the initial mapping sync has completed.
func (c *Checker) SetInitialSyncComplete() {
c.mu.Lock()
defer c.mu.Unlock()
c.initialSyncComplete = true
}
// SetShuttingDown marks the server as shutting down.
// This causes ReadinessProbe to return false so load balancers stop routing traffic.
func (c *Checker) SetShuttingDown() {
c.mu.Lock()
defer c.mu.Unlock()
c.shuttingDown = true
}
// CheckClientsConnected verifies all clients are connected to management/signal/relay.
// Uses the provided context for timeout/cancellation, with a maximum bound of maxClientCheckTimeout.
// Limits concurrent checks via semaphore.
func (c *Checker) CheckClientsConnected(ctx context.Context) (bool, map[types.AccountID]ClientHealth) {
// Apply upper bound timeout in case parent context has no deadline
ctx, cancel := context.WithTimeout(ctx, maxClientCheckTimeout)
defer cancel()
clients := c.provider.ListClientsForStartup()
// No clients is not a health issue
if len(clients) == 0 {
return true, make(map[types.AccountID]ClientHealth)
}
type result struct {
accountID types.AccountID
health ClientHealth
}
resultsCh := make(chan result, len(clients))
var wg sync.WaitGroup
for accountID, client := range clients {
wg.Add(1)
go func(id types.AccountID, cl *embed.Client) {
defer wg.Done()
// Acquire semaphore
select {
case c.checkSem <- struct{}{}:
defer func() { <-c.checkSem }()
case <-ctx.Done():
resultsCh <- result{id, ClientHealth{Healthy: false, Error: ctx.Err().Error()}}
return
}
resultsCh <- result{id, c.checkHealth(cl)}
}(accountID, client)
}
go func() {
wg.Wait()
close(resultsCh)
}()
results := make(map[types.AccountID]ClientHealth)
allHealthy := true
for r := range resultsCh {
results[r.accountID] = r.health
if !r.health.Healthy {
allHealthy = false
}
}
return allHealthy, results
}
// LivenessProbe returns true if the process is alive.
// This should always return true if we can respond.
func (c *Checker) LivenessProbe() bool {
return true
}
// ReadinessProbe returns true if the server can accept traffic.
func (c *Checker) ReadinessProbe() bool {
c.mu.RLock()
defer c.mu.RUnlock()
if c.shuttingDown {
return false
}
return c.managementConnected
}
// StartupProbe checks if initial startup is complete.
// Checks management connection, initial sync, and all client health directly.
// Uses the provided context for timeout/cancellation.
func (c *Checker) StartupProbe(ctx context.Context) bool {
c.mu.RLock()
mgmt := c.managementConnected
sync := c.initialSyncComplete
c.mu.RUnlock()
if !mgmt || !sync {
return false
}
// Check all clients are connected to management/signal/relay.
// Returns true when no clients exist (nothing to check).
allHealthy, _ := c.CheckClientsConnected(ctx)
return allHealthy
}
// Handler returns an http.Handler for health probe endpoints.
func (c *Checker) Handler() http.Handler {
mux := http.NewServeMux()
mux.HandleFunc("/healthz/live", c.handleLiveness)
mux.HandleFunc("/healthz/ready", c.handleReadiness)
mux.HandleFunc("/healthz/startup", c.handleStartup)
mux.HandleFunc("/healthz", c.handleFull)
return mux
}
func (c *Checker) handleLiveness(w http.ResponseWriter, r *http.Request) {
if c.LivenessProbe() {
c.writeProbeResponse(w, http.StatusOK, "ok", nil, nil)
return
}
c.writeProbeResponse(w, http.StatusServiceUnavailable, "fail", nil, nil)
}
func (c *Checker) handleReadiness(w http.ResponseWriter, r *http.Request) {
c.mu.RLock()
checks := map[string]bool{
"management_connected": c.managementConnected,
}
c.mu.RUnlock()
if c.ReadinessProbe() {
c.writeProbeResponse(w, http.StatusOK, "ok", checks, nil)
return
}
c.writeProbeResponse(w, http.StatusServiceUnavailable, "fail", checks, nil)
}
func (c *Checker) handleStartup(w http.ResponseWriter, r *http.Request) {
c.mu.RLock()
mgmt := c.managementConnected
syncComplete := c.initialSyncComplete
c.mu.RUnlock()
allClientsHealthy, clientHealth := c.CheckClientsConnected(r.Context())
checks := map[string]bool{
"management_connected": mgmt,
"initial_sync_complete": syncComplete,
"all_clients_healthy": allClientsHealthy,
}
ready := mgmt && syncComplete && allClientsHealthy
if ready {
c.writeProbeResponse(w, http.StatusOK, "ok", checks, clientHealth)
return
}
c.writeProbeResponse(w, http.StatusServiceUnavailable, "fail", checks, clientHealth)
}
func (c *Checker) handleFull(w http.ResponseWriter, r *http.Request) {
c.mu.RLock()
mgmt := c.managementConnected
sync := c.initialSyncComplete
c.mu.RUnlock()
allClientsHealthy, clientHealth := c.CheckClientsConnected(r.Context())
checks := map[string]bool{
"management_connected": mgmt,
"initial_sync_complete": sync,
"all_clients_healthy": allClientsHealthy,
}
status := "ok"
statusCode := http.StatusOK
if !c.ReadinessProbe() {
status = "fail"
statusCode = http.StatusServiceUnavailable
}
c.writeProbeResponse(w, statusCode, status, checks, clientHealth)
}
func (c *Checker) writeProbeResponse(w http.ResponseWriter, statusCode int, status string, checks map[string]bool, clients map[types.AccountID]ClientHealth) {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(statusCode)
resp := ProbeResponse{
Status: status,
Checks: checks,
Clients: clients,
}
if err := json.NewEncoder(w).Encode(resp); err != nil {
c.logger.Debugf("write health response: %v", err)
}
}
// ListenAndServe starts the health probe server.
func (s *Server) ListenAndServe() error {
s.logger.Infof("starting health probe server on %s", s.server.Addr)
return s.server.ListenAndServe()
}
// Serve starts the health probe server on the given listener.
func (s *Server) Serve(l net.Listener) error {
s.logger.Infof("starting health probe server on %s", l.Addr())
return s.server.Serve(l)
}
// Shutdown gracefully shuts down the health probe server.
func (s *Server) Shutdown(ctx context.Context) error {
return s.server.Shutdown(ctx)
}
// NewChecker creates a new health checker.
func NewChecker(logger *log.Logger, provider clientProvider) *Checker {
if logger == nil {
logger = log.StandardLogger()
}
return &Checker{
logger: logger,
provider: provider,
checkSem: make(chan struct{}, maxConcurrentChecks),
checkHealth: checkClientHealth,
}
}
// NewServer creates a new health probe server.
// If metricsHandler is non-nil, it is mounted at /metrics on the same port.
func NewServer(addr string, checker *Checker, logger *log.Logger, metricsHandler http.Handler) *Server {
if logger == nil {
logger = log.StandardLogger()
}
handler := checker.Handler()
if metricsHandler != nil {
mux := http.NewServeMux()
mux.Handle("/metrics", metricsHandler)
mux.Handle("/", handler)
handler = mux
}
return &Server{
server: &http.Server{
Addr: addr,
Handler: handler,
ReadTimeout: 5 * time.Second,
WriteTimeout: 5 * time.Second,
},
logger: logger,
checker: checker,
}
}
func checkClientHealth(client *embed.Client) ClientHealth {
if client == nil {
return ClientHealth{
Healthy: false,
Error: "client not initialized",
}
}
status, err := client.Status()
if err != nil {
return ClientHealth{
Healthy: false,
Error: err.Error(),
}
}
// Count only rel:// and rels:// relays (not stun/turn)
var relayCount, relaysConnected int
for _, relay := range status.Relays {
if !strings.HasPrefix(relay.URI, "rel://") && !strings.HasPrefix(relay.URI, "rels://") {
continue
}
relayCount++
if relay.Err == nil {
relaysConnected++
}
}
// Count peer connection stats
now := time.Now()
var peersConnected, peersP2P, peersRelayed, peersDegraded int
for _, p := range status.Peers {
if p.ConnStatus != embed.PeerStatusConnected {
continue
}
peersConnected++
if p.Relayed {
peersRelayed++
} else {
peersP2P++
}
if p.LastWireguardHandshake.IsZero() || now.Sub(p.LastWireguardHandshake) > handshakeStaleThreshold {
peersDegraded++
}
}
// Client is healthy if connected to management, signal, and at least one relay (if any are defined)
healthy := status.ManagementState.Connected &&
status.SignalState.Connected &&
(relayCount == 0 || relaysConnected > 0)
return ClientHealth{
Healthy: healthy,
ManagementConnected: status.ManagementState.Connected,
SignalConnected: status.SignalState.Connected,
RelaysConnected: relaysConnected,
RelaysTotal: relayCount,
PeersTotal: len(status.Peers),
PeersConnected: peersConnected,
PeersP2P: peersP2P,
PeersRelayed: peersRelayed,
PeersDegraded: peersDegraded,
}
}
proxy/internal/health/health_test.go
+473
View File
@@ -0,0 +1,473 @@
package health
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"sync"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/netbirdio/netbird/client/embed"
"github.com/netbirdio/netbird/proxy/internal/types"
)
type mockClientProvider struct {
clients map[types.AccountID]*embed.Client
}
func (m *mockClientProvider) ListClientsForStartup() map[types.AccountID]*embed.Client {
return m.clients
}
// newTestChecker creates a checker with a mock health function for testing.
// The health function returns the provided ClientHealth for every client.
func newTestChecker(provider clientProvider, healthResult ClientHealth) *Checker {
c := NewChecker(nil, provider)
c.checkHealth = func(_ *embed.Client) ClientHealth {
return healthResult
}
return c
}
func TestChecker_LivenessProbe(t *testing.T) {
checker := NewChecker(nil, &mockClientProvider{})
// Liveness should always return true if we can respond.
assert.True(t, checker.LivenessProbe())
}
func TestChecker_ReadinessProbe(t *testing.T) {
checker := NewChecker(nil, &mockClientProvider{})
// Initially not ready (management not connected).
assert.False(t, checker.ReadinessProbe())
// After management connects, should be ready.
checker.SetManagementConnected(true)
assert.True(t, checker.ReadinessProbe())
// If management disconnects, should not be ready.
checker.SetManagementConnected(false)
assert.False(t, checker.ReadinessProbe())
}
// TestStartupProbe_EmptyServiceList covers the scenario where management has
// no services configured for this proxy. The proxy should become ready once
// management is connected and the initial sync completes, even with zero clients.
func TestStartupProbe_EmptyServiceList(t *testing.T) {
checker := NewChecker(nil, &mockClientProvider{})
// No management connection = not ready.
assert.False(t, checker.StartupProbe(context.Background()))
// Management connected but no sync = not ready.
checker.SetManagementConnected(true)
assert.False(t, checker.StartupProbe(context.Background()))
// Management + sync complete + no clients = ready.
checker.SetInitialSyncComplete()
assert.True(t, checker.StartupProbe(context.Background()))
}
// TestStartupProbe_WithUnhealthyClients verifies that when services exist
// and clients have been created but are not yet fully connected (to mgmt,
// signal, relays), the startup probe does NOT pass.
func TestStartupProbe_WithUnhealthyClients(t *testing.T) {
provider := &mockClientProvider{
clients: map[types.AccountID]*embed.Client{
"account-1": nil, // concrete client not needed; checkHealth is mocked
"account-2": nil,
},
}
checker := newTestChecker(provider, ClientHealth{Healthy: false, Error: "not connected yet"})
checker.SetManagementConnected(true)
checker.SetInitialSyncComplete()
assert.False(t, checker.StartupProbe(context.Background()),
"startup probe must not pass when clients are unhealthy")
}
// TestStartupProbe_WithHealthyClients verifies that once all clients are
// connected and healthy, the startup probe passes.
func TestStartupProbe_WithHealthyClients(t *testing.T) {
provider := &mockClientProvider{
clients: map[types.AccountID]*embed.Client{
"account-1": nil,
"account-2": nil,
},
}
checker := newTestChecker(provider, ClientHealth{
Healthy: true,
ManagementConnected: true,
SignalConnected: true,
RelaysConnected: 1,
RelaysTotal: 1,
})
checker.SetManagementConnected(true)
checker.SetInitialSyncComplete()
assert.True(t, checker.StartupProbe(context.Background()),
"startup probe must pass when all clients are healthy")
}
// TestStartupProbe_MixedHealthClients verifies that if any single client is
// unhealthy, the startup probe fails (all-or-nothing).
func TestStartupProbe_MixedHealthClients(t *testing.T) {
provider := &mockClientProvider{
clients: map[types.AccountID]*embed.Client{
"healthy-account": nil,
"unhealthy-account": nil,
},
}
checker := NewChecker(nil, provider)
checker.checkHealth = func(cl *embed.Client) ClientHealth {
// We identify accounts by their position in the map iteration; since we
// can't control map order, make exactly one unhealthy via counter.
return ClientHealth{Healthy: false}
}
checker.SetManagementConnected(true)
checker.SetInitialSyncComplete()
assert.False(t, checker.StartupProbe(context.Background()),
"startup probe must fail if any client is unhealthy")
}
// TestStartupProbe_RequiresAllConditions ensures that each individual
// prerequisite (management, sync, clients) is necessary. The probe must not
// pass if any one is missing.
func TestStartupProbe_RequiresAllConditions(t *testing.T) {
provider := &mockClientProvider{
clients: map[types.AccountID]*embed.Client{
"account-1": nil,
},
}
t.Run("no management", func(t *testing.T) {
checker := newTestChecker(provider, ClientHealth{Healthy: true})
checker.SetInitialSyncComplete()
// management NOT connected
assert.False(t, checker.StartupProbe(context.Background()))
})
t.Run("no sync", func(t *testing.T) {
checker := newTestChecker(provider, ClientHealth{Healthy: true})
checker.SetManagementConnected(true)
// sync NOT complete
assert.False(t, checker.StartupProbe(context.Background()))
})
t.Run("unhealthy client", func(t *testing.T) {
checker := newTestChecker(provider, ClientHealth{Healthy: false})
checker.SetManagementConnected(true)
checker.SetInitialSyncComplete()
assert.False(t, checker.StartupProbe(context.Background()))
})
t.Run("all conditions met", func(t *testing.T) {
checker := newTestChecker(provider, ClientHealth{Healthy: true})
checker.SetManagementConnected(true)
checker.SetInitialSyncComplete()
assert.True(t, checker.StartupProbe(context.Background()))
})
}
// TestStartupProbe_ConcurrentAccess runs the startup probe from many
// goroutines simultaneously to check for races.
func TestStartupProbe_ConcurrentAccess(t *testing.T) {
provider := &mockClientProvider{
clients: map[types.AccountID]*embed.Client{
"account-1": nil,
"account-2": nil,
},
}
checker := newTestChecker(provider, ClientHealth{Healthy: true})
checker.SetManagementConnected(true)
checker.SetInitialSyncComplete()
var wg sync.WaitGroup
const goroutines = 50
results := make([]bool, goroutines)
for i := range goroutines {
wg.Add(1)
go func(idx int) {
defer wg.Done()
results[idx] = checker.StartupProbe(context.Background())
}(i)
}
wg.Wait()
for i, r := range results {
assert.True(t, r, "goroutine %d got unexpected result", i)
}
}
// TestStartupProbe_CancelledContext verifies that a cancelled context causes
// the probe to report unhealthy when client checks are needed.
func TestStartupProbe_CancelledContext(t *testing.T) {
t.Run("no management bypasses context", func(t *testing.T) {
checker := NewChecker(nil, &mockClientProvider{})
ctx, cancel := context.WithCancel(context.Background())
cancel()
// Should be false because management isn't connected, context is irrelevant.
assert.False(t, checker.StartupProbe(ctx))
})
t.Run("with clients and cancelled context", func(t *testing.T) {
provider := &mockClientProvider{
clients: map[types.AccountID]*embed.Client{
"account-1": nil,
},
}
checker := NewChecker(nil, provider)
// Use the real checkHealth path — a cancelled context should cause
// the semaphore acquisition to fail, reporting unhealthy.
checker.SetManagementConnected(true)
checker.SetInitialSyncComplete()
ctx, cancel := context.WithCancel(context.Background())
cancel()
assert.False(t, checker.StartupProbe(ctx),
"cancelled context must result in unhealthy when clients exist")
})
}
// TestHandler_Startup_EmptyServiceList verifies the HTTP startup endpoint
// returns 200 when management is connected, sync is complete, and there are
// no services/clients.
func TestHandler_Startup_EmptyServiceList(t *testing.T) {
checker := NewChecker(nil, &mockClientProvider{})
checker.SetManagementConnected(true)
checker.SetInitialSyncComplete()
handler := checker.Handler()
req := httptest.NewRequest(http.MethodGet, "/healthz/startup", nil)
rec := httptest.NewRecorder()
handler.ServeHTTP(rec, req)
assert.Equal(t, http.StatusOK, rec.Code)
var resp ProbeResponse
require.NoError(t, json.NewDecoder(rec.Body).Decode(&resp))
assert.Equal(t, "ok", resp.Status)
assert.True(t, resp.Checks["management_connected"])
assert.True(t, resp.Checks["initial_sync_complete"])
assert.True(t, resp.Checks["all_clients_healthy"])
assert.Empty(t, resp.Clients)
}
// TestHandler_Startup_WithUnhealthyClients verifies that the HTTP startup
// endpoint returns 503 when clients exist but are not yet healthy.
func TestHandler_Startup_WithUnhealthyClients(t *testing.T) {
provider := &mockClientProvider{
clients: map[types.AccountID]*embed.Client{
"account-1": nil,
},
}
checker := newTestChecker(provider, ClientHealth{Healthy: false, Error: "starting"})
checker.SetManagementConnected(true)
checker.SetInitialSyncComplete()
handler := checker.Handler()
req := httptest.NewRequest(http.MethodGet, "/healthz/startup", nil)
rec := httptest.NewRecorder()
handler.ServeHTTP(rec, req)
assert.Equal(t, http.StatusServiceUnavailable, rec.Code)
var resp ProbeResponse
require.NoError(t, json.NewDecoder(rec.Body).Decode(&resp))
assert.Equal(t, "fail", resp.Status)
assert.True(t, resp.Checks["management_connected"])
assert.True(t, resp.Checks["initial_sync_complete"])
assert.False(t, resp.Checks["all_clients_healthy"])
require.Contains(t, resp.Clients, types.AccountID("account-1"))
assert.Equal(t, "starting", resp.Clients["account-1"].Error)
}
// TestHandler_Startup_WithHealthyClients verifies the HTTP startup endpoint
// returns 200 once clients are healthy.
func TestHandler_Startup_WithHealthyClients(t *testing.T) {
provider := &mockClientProvider{
clients: map[types.AccountID]*embed.Client{
"account-1": nil,
},
}
checker := newTestChecker(provider, ClientHealth{
Healthy: true,
ManagementConnected: true,
SignalConnected: true,
RelaysConnected: 1,
RelaysTotal: 1,
})
checker.SetManagementConnected(true)
checker.SetInitialSyncComplete()
handler := checker.Handler()
req := httptest.NewRequest(http.MethodGet, "/healthz/startup", nil)
rec := httptest.NewRecorder()
handler.ServeHTTP(rec, req)
assert.Equal(t, http.StatusOK, rec.Code)
var resp ProbeResponse
require.NoError(t, json.NewDecoder(rec.Body).Decode(&resp))
assert.Equal(t, "ok", resp.Status)
assert.True(t, resp.Checks["all_clients_healthy"])
}
// TestHandler_Startup_NotComplete verifies the startup handler returns 503
// when prerequisites aren't met.
func TestHandler_Startup_NotComplete(t *testing.T) {
checker := NewChecker(nil, &mockClientProvider{})
handler := checker.Handler()
req := httptest.NewRequest(http.MethodGet, "/healthz/startup", nil)
rec := httptest.NewRecorder()
handler.ServeHTTP(rec, req)
assert.Equal(t, http.StatusServiceUnavailable, rec.Code)
var resp ProbeResponse
require.NoError(t, json.NewDecoder(rec.Body).Decode(&resp))
assert.Equal(t, "fail", resp.Status)
}
func TestChecker_Handler_Liveness(t *testing.T) {
checker := NewChecker(nil, &mockClientProvider{})
handler := checker.Handler()
req := httptest.NewRequest(http.MethodGet, "/healthz/live", nil)
rec := httptest.NewRecorder()
handler.ServeHTTP(rec, req)
assert.Equal(t, http.StatusOK, rec.Code)
var resp ProbeResponse
require.NoError(t, json.NewDecoder(rec.Body).Decode(&resp))
assert.Equal(t, "ok", resp.Status)
}
func TestChecker_Handler_Readiness_NotReady(t *testing.T) {
checker := NewChecker(nil, &mockClientProvider{})
handler := checker.Handler()
req := httptest.NewRequest(http.MethodGet, "/healthz/ready", nil)
rec := httptest.NewRecorder()
handler.ServeHTTP(rec, req)
assert.Equal(t, http.StatusServiceUnavailable, rec.Code)
var resp ProbeResponse
require.NoError(t, json.NewDecoder(rec.Body).Decode(&resp))
assert.Equal(t, "fail", resp.Status)
assert.False(t, resp.Checks["management_connected"])
}
func TestChecker_Handler_Readiness_Ready(t *testing.T) {
checker := NewChecker(nil, &mockClientProvider{})
checker.SetManagementConnected(true)
handler := checker.Handler()
req := httptest.NewRequest(http.MethodGet, "/healthz/ready", nil)
rec := httptest.NewRecorder()
handler.ServeHTTP(rec, req)
assert.Equal(t, http.StatusOK, rec.Code)
var resp ProbeResponse
require.NoError(t, json.NewDecoder(rec.Body).Decode(&resp))
assert.Equal(t, "ok", resp.Status)
assert.True(t, resp.Checks["management_connected"])
}
func TestChecker_Handler_Full(t *testing.T) {
checker := NewChecker(nil, &mockClientProvider{})
checker.SetManagementConnected(true)
handler := checker.Handler()
req := httptest.NewRequest(http.MethodGet, "/healthz", nil)
rec := httptest.NewRecorder()
handler.ServeHTTP(rec, req)
assert.Equal(t, http.StatusOK, rec.Code)
var resp ProbeResponse
require.NoError(t, json.NewDecoder(rec.Body).Decode(&resp))
assert.Equal(t, "ok", resp.Status)
assert.NotNil(t, resp.Checks)
// Clients may be empty map when no clients exist.
assert.Empty(t, resp.Clients)
}
func TestChecker_SetShuttingDown(t *testing.T) {
checker := NewChecker(nil, &mockClientProvider{})
checker.SetManagementConnected(true)
assert.True(t, checker.ReadinessProbe(), "should be ready before shutdown")
checker.SetShuttingDown()
assert.False(t, checker.ReadinessProbe(), "should not be ready after shutdown")
}
func TestChecker_Handler_Readiness_ShuttingDown(t *testing.T) {
checker := NewChecker(nil, &mockClientProvider{})
checker.SetManagementConnected(true)
checker.SetShuttingDown()
handler := checker.Handler()
req := httptest.NewRequest(http.MethodGet, "/healthz/ready", nil)
rec := httptest.NewRecorder()
handler.ServeHTTP(rec, req)
assert.Equal(t, http.StatusServiceUnavailable, rec.Code)
var resp ProbeResponse
require.NoError(t, json.NewDecoder(rec.Body).Decode(&resp))
assert.Equal(t, "fail", resp.Status)
}
func TestNewServer_WithMetricsHandler(t *testing.T) {
checker := NewChecker(nil, &mockClientProvider{})
checker.SetManagementConnected(true)
metricsHandler := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte("metrics"))
})
srv := NewServer(":0", checker, nil, metricsHandler)
require.NotNil(t, srv)
// Verify health endpoint still works through the mux.
req := httptest.NewRequest(http.MethodGet, "/healthz/live", nil)
rec := httptest.NewRecorder()
srv.server.Handler.ServeHTTP(rec, req)
assert.Equal(t, http.StatusOK, rec.Code)
// Verify metrics endpoint is mounted.
req = httptest.NewRequest(http.MethodGet, "/metrics", nil)
rec = httptest.NewRecorder()
srv.server.Handler.ServeHTTP(rec, req)
assert.Equal(t, http.StatusOK, rec.Code)
assert.Equal(t, "metrics", rec.Body.String())
}
func TestNewServer_WithoutMetricsHandler(t *testing.T) {
checker := NewChecker(nil, &mockClientProvider{})
checker.SetManagementConnected(true)
srv := NewServer(":0", checker, nil, nil)
require.NotNil(t, srv)
req := httptest.NewRequest(http.MethodGet, "/healthz/live", nil)
rec := httptest.NewRecorder()
srv.server.Handler.ServeHTTP(rec, req)
assert.Equal(t, http.StatusOK, rec.Code)
}
proxy/internal/k8s/lease.go
+281
View File
@@ -0,0 +1,281 @@
// Package k8s provides a lightweight Kubernetes API client for coordination
// Leases. It uses raw HTTP calls against the mounted service account
// credentials, avoiding a dependency on client-go.
package k8s
import (
"bytes"
"context"
"crypto/sha256"
"crypto/tls"
"crypto/x509"
"encoding/hex"
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
"os"
"strings"
"time"
)
const (
saTokenPath = "/var/run/secrets/kubernetes.io/serviceaccount/token" //nolint:gosec
saNamespacePath = "/var/run/secrets/kubernetes.io/serviceaccount/namespace"
saCACertPath = "/var/run/secrets/kubernetes.io/serviceaccount/ca.crt"
leaseAPIPath = "/apis/coordination.k8s.io/v1"
)
// ErrConflict is returned when a Lease update fails due to a
// resourceVersion mismatch (another writer updated the object first).
var ErrConflict = errors.New("conflict: resource version mismatch")
// Lease represents a coordination.k8s.io/v1 Lease object with only the
// fields needed for distributed locking.
type Lease struct {
APIVersion string `json:"apiVersion"`
Kind string `json:"kind"`
Metadata LeaseMetadata `json:"metadata"`
Spec LeaseSpec `json:"spec"`
}
// LeaseMetadata holds the standard k8s object metadata fields used by Leases.
type LeaseMetadata struct {
Name string `json:"name"`
Namespace string `json:"namespace,omitempty"`
ResourceVersion string `json:"resourceVersion,omitempty"`
Annotations map[string]string `json:"annotations,omitempty"`
}
// LeaseSpec holds the Lease specification fields.
type LeaseSpec struct {
HolderIdentity *string `json:"holderIdentity"`
LeaseDurationSeconds *int32 `json:"leaseDurationSeconds,omitempty"`
AcquireTime *MicroTime `json:"acquireTime"`
RenewTime *MicroTime `json:"renewTime"`
}
// MicroTime wraps time.Time with Kubernetes MicroTime JSON formatting.
type MicroTime struct {
time.Time
}
const microTimeFormat = "2006-01-02T15:04:05.000000Z"
// MarshalJSON implements json.Marshaler with k8s MicroTime format.
func (t *MicroTime) MarshalJSON() ([]byte, error) {
return json.Marshal(t.UTC().Format(microTimeFormat))
}
// UnmarshalJSON implements json.Unmarshaler with k8s MicroTime format.
func (t *MicroTime) UnmarshalJSON(data []byte) error {
var s string
if err := json.Unmarshal(data, &s); err != nil {
return err
}
if s == "" {
t.Time = time.Time{}
return nil
}
parsed, err := time.Parse(microTimeFormat, s)
if err != nil {
return fmt.Errorf("parse MicroTime %q: %w", s, err)
}
t.Time = parsed
return nil
}
// LeaseClient talks to the Kubernetes coordination API using raw HTTP.
type LeaseClient struct {
baseURL string
namespace string
httpClient *http.Client
}
// NewLeaseClient creates a client that authenticates via the pod's
// mounted service account. It reads the namespace and CA certificate
// at construction time (they don't rotate) but reads the bearer token
// fresh on each request (tokens rotate).
func NewLeaseClient() (*LeaseClient, error) {
host := os.Getenv("KUBERNETES_SERVICE_HOST")
port := os.Getenv("KUBERNETES_SERVICE_PORT")
if host == "" || port == "" {
return nil, fmt.Errorf("KUBERNETES_SERVICE_HOST/PORT not set")
}
ns, err := os.ReadFile(saNamespacePath)
if err != nil {
return nil, fmt.Errorf("read namespace from %s: %w", saNamespacePath, err)
}
caCert, err := os.ReadFile(saCACertPath)
if err != nil {
return nil, fmt.Errorf("read CA cert from %s: %w", saCACertPath, err)
}
pool := x509.NewCertPool()
if !pool.AppendCertsFromPEM(caCert) {
return nil, fmt.Errorf("parse CA certificate from %s", saCACertPath)
}
return &LeaseClient{
baseURL: fmt.Sprintf("https://%s:%s", host, port),
namespace: strings.TrimSpace(string(ns)),
httpClient: &http.Client{
Timeout: 10 * time.Second,
Transport: &http.Transport{
TLSClientConfig: &tls.Config{
RootCAs: pool,
},
},
},
}, nil
}
// Namespace returns the namespace this client operates in.
func (c *LeaseClient) Namespace() string {
return c.namespace
}
// Get retrieves a Lease by name. Returns (nil, nil) if the Lease does not exist.
func (c *LeaseClient) Get(ctx context.Context, name string) (*Lease, error) {
url := fmt.Sprintf("%s%s/namespaces/%s/leases/%s", c.baseURL, leaseAPIPath, c.namespace, name)
resp, err := c.doRequest(ctx, http.MethodGet, url, nil)
if err != nil {
return nil, err
}
defer func() { _ = resp.Body.Close() }()
if resp.StatusCode == http.StatusNotFound {
return nil, nil //nolint:nilnil
}
if resp.StatusCode != http.StatusOK {
return nil, c.readError(resp)
}
var lease Lease
if err := json.NewDecoder(resp.Body).Decode(&lease); err != nil {
return nil, fmt.Errorf("decode lease response: %w", err)
}
return &lease, nil
}
// Create creates a new Lease. Returns the created Lease with server-assigned
// fields like resourceVersion populated.
func (c *LeaseClient) Create(ctx context.Context, lease *Lease) (*Lease, error) {
url := fmt.Sprintf("%s%s/namespaces/%s/leases", c.baseURL, leaseAPIPath, c.namespace)
lease.APIVersion = "coordination.k8s.io/v1"
lease.Kind = "Lease"
if lease.Metadata.Namespace == "" {
lease.Metadata.Namespace = c.namespace
}
resp, err := c.doRequest(ctx, http.MethodPost, url, lease)
if err != nil {
return nil, err
}
defer func() { _ = resp.Body.Close() }()
if resp.StatusCode == http.StatusConflict {
return nil, ErrConflict
}
if resp.StatusCode != http.StatusCreated {
return nil, c.readError(resp)
}
var created Lease
if err := json.NewDecoder(resp.Body).Decode(&created); err != nil {
return nil, fmt.Errorf("decode created lease: %w", err)
}
return &created, nil
}
// Update replaces a Lease. The lease.Metadata.ResourceVersion must match
// the current server value (optimistic concurrency). Returns ErrConflict
// on version mismatch.
func (c *LeaseClient) Update(ctx context.Context, lease *Lease) (*Lease, error) {
url := fmt.Sprintf("%s%s/namespaces/%s/leases/%s", c.baseURL, leaseAPIPath, c.namespace, lease.Metadata.Name)
lease.APIVersion = "coordination.k8s.io/v1"
lease.Kind = "Lease"
resp, err := c.doRequest(ctx, http.MethodPut, url, lease)
if err != nil {
return nil, err
}
defer func() { _ = resp.Body.Close() }()
if resp.StatusCode == http.StatusConflict {
return nil, ErrConflict
}
if resp.StatusCode != http.StatusOK {
return nil, c.readError(resp)
}
var updated Lease
if err := json.NewDecoder(resp.Body).Decode(&updated); err != nil {
return nil, fmt.Errorf("decode updated lease: %w", err)
}
return &updated, nil
}
func (c *LeaseClient) doRequest(ctx context.Context, method, url string, body any) (*http.Response, error) {
token, err := readToken()
if err != nil {
return nil, fmt.Errorf("read service account token: %w", err)
}
var bodyReader io.Reader
if body != nil {
data, err := json.Marshal(body)
if err != nil {
return nil, fmt.Errorf("marshal request body: %w", err)
}
bodyReader = bytes.NewReader(data)
}
req, err := http.NewRequestWithContext(ctx, method, url, bodyReader)
if err != nil {
return nil, fmt.Errorf("create request: %w", err)
}
req.Header.Set("Authorization", "Bearer "+token)
req.Header.Set("Accept", "application/json")
if body != nil {
req.Header.Set("Content-Type", "application/json")
}
return c.httpClient.Do(req)
}
func readToken() (string, error) {
data, err := os.ReadFile(saTokenPath)
if err != nil {
return "", fmt.Errorf("read %s: %w", saTokenPath, err)
}
return strings.TrimSpace(string(data)), nil
}
func (c *LeaseClient) readError(resp *http.Response) error {
body, _ := io.ReadAll(io.LimitReader(resp.Body, 1024))
return fmt.Errorf("k8s API %s %d: %s", resp.Request.URL.Path, resp.StatusCode, string(body))
}
// LeaseNameForDomain returns a deterministic, DNS-label-safe Lease name
// for the given domain. The domain is hashed to avoid dots and length issues.
func LeaseNameForDomain(domain string) string {
h := sha256.Sum256([]byte(domain))
return "cert-lock-" + hex.EncodeToString(h[:8])
}
// InCluster reports whether the process is running inside a Kubernetes pod
// by checking for the KUBERNETES_SERVICE_HOST environment variable.
func InCluster() bool {
_, exists := os.LookupEnv("KUBERNETES_SERVICE_HOST")
return exists
}
proxy/internal/k8s/lease_test.go
+102
View File
@@ -0,0 +1,102 @@
package k8s
import (
"encoding/json"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestLeaseNameForDomain(t *testing.T) {
tests := []struct {
domain string
}{
{"example.com"},
{"app.example.com"},
{"another.domain.io"},
}
seen := make(map[string]string)
for _, tc := range tests {
name := LeaseNameForDomain(tc.domain)
assert.True(t, len(name) <= 63, "must be valid DNS label length")
assert.Regexp(t, `^cert-lock-[0-9a-f]{16}$`, name,
"must match expected format for domain %q", tc.domain)
// Same input produces same output.
assert.Equal(t, name, LeaseNameForDomain(tc.domain), "must be deterministic")
// Different domains produce different names.
if prev, ok := seen[name]; ok {
t.Errorf("collision: %q and %q both map to %s", prev, tc.domain, name)
}
seen[name] = tc.domain
}
}
func TestMicroTimeJSON(t *testing.T) {
ts := time.Date(2024, 6, 15, 10, 30, 0, 0, time.UTC)
mt := &MicroTime{Time: ts}
data, err := json.Marshal(mt)
require.NoError(t, err)
assert.Equal(t, `"2024-06-15T10:30:00.000000Z"`, string(data))
var decoded MicroTime
require.NoError(t, json.Unmarshal(data, &decoded))
assert.True(t, ts.Equal(decoded.Time), "round-trip should preserve time")
}
func TestMicroTimeNullJSON(t *testing.T) {
// Null pointer serializes as JSON null via the Lease struct.
spec := LeaseSpec{
HolderIdentity: nil,
AcquireTime: nil,
RenewTime: nil,
}
data, err := json.Marshal(spec)
require.NoError(t, err)
assert.Contains(t, string(data), `"acquireTime":null`)
assert.Contains(t, string(data), `"renewTime":null`)
}
func TestLeaseJSONRoundTrip(t *testing.T) {
holder := "pod-abc"
dur := int32(300)
now := MicroTime{Time: time.Now().UTC().Truncate(time.Microsecond)}
original := Lease{
APIVersion: "coordination.k8s.io/v1",
Kind: "Lease",
Metadata: LeaseMetadata{
Name: "cert-lock-abcdef0123456789",
Namespace: "default",
ResourceVersion: "12345",
Annotations: map[string]string{
"netbird.io/domain": "app.example.com",
},
},
Spec: LeaseSpec{
HolderIdentity: &holder,
LeaseDurationSeconds: &dur,
AcquireTime: &now,
RenewTime: &now,
},
}
data, err := json.Marshal(original)
require.NoError(t, err)
var decoded Lease
require.NoError(t, json.Unmarshal(data, &decoded))
assert.Equal(t, original.Metadata.Name, decoded.Metadata.Name)
assert.Equal(t, original.Metadata.ResourceVersion, decoded.Metadata.ResourceVersion)
assert.Equal(t, *original.Spec.HolderIdentity, *decoded.Spec.HolderIdentity)
assert.Equal(t, *original.Spec.LeaseDurationSeconds, *decoded.Spec.LeaseDurationSeconds)
assert.True(t, original.Spec.AcquireTime.Equal(decoded.Spec.AcquireTime.Time))
}
proxy/internal/llm/anthropic.go
+196
View File
@@ -0,0 +1,196 @@
package llm
import (
"encoding/json"
"fmt"
"strings"
)
// AnthropicParser implements the Parser interface for the Anthropic Messages
// and Completions APIs. Detection is substring-based to tolerate upstream
// path rewrites.
type AnthropicParser struct{}
var anthropicPathHints = []string{
"/v1/messages",
"/v1/complete",
}
// Provider returns ProviderAnthropic.
func (AnthropicParser) Provider() Provider { return ProviderAnthropic }
// ProviderName returns the stable label used for metrics and metadata.
func (AnthropicParser) ProviderName() string { return "anthropic" }
// DetectFromURL reports whether the given request path looks like an
// Anthropic API endpoint. The match is case-insensitive and substring-based.
func (AnthropicParser) DetectFromURL(path string) bool {
lower := strings.ToLower(path)
for _, hint := range anthropicPathHints {
if strings.Contains(lower, hint) {
return true
}
}
return false
}
type anthropicRequest struct {
Model string `json:"model"`
Stream *bool `json:"stream"`
System json.RawMessage `json:"system"`
Messages []anthropicMessage `json:"messages"`
// Legacy /v1/complete endpoint.
Prompt string `json:"prompt"`
}
type anthropicMessage struct {
Role string `json:"role"`
Content json.RawMessage `json:"content"`
}
// ParseRequest extracts the model name and streaming flag from an Anthropic
// request body. Unknown or missing fields leave the corresponding struct
// members zero-valued.
func (AnthropicParser) ParseRequest(body []byte) (RequestFacts, error) {
var req anthropicRequest
if err := json.Unmarshal(body, &req); err != nil {
return RequestFacts{}, fmt.Errorf("decode anthropic request: %w: %v", ErrMalformedRequest, err)
}
return RequestFacts{
Model: req.Model,
Stream: ptrDeref(req.Stream),
}, nil
}
type anthropicResponse struct {
Usage struct {
InputTokens int64 `json:"input_tokens"`
OutputTokens int64 `json:"output_tokens"`
// CacheReadInputTokens and CacheCreationInputTokens are
// ADDITIVE to InputTokens (not subset), each billed at its
// own rate by the cost meter. cache_read is the cheaper
// read-from-cache rate, cache_creation is the more
// expensive write-to-cache rate.
CacheReadInputTokens int64 `json:"cache_read_input_tokens"`
CacheCreationInputTokens int64 `json:"cache_creation_input_tokens"`
} `json:"usage"`
}
// ParseResponse decodes the non-streaming Anthropic response envelope. Status
// codes other than 200 are treated as non-LLM responses so the caller can
// skip cost accounting without aborting the request.
func (AnthropicParser) ParseResponse(status int, contentType string, body []byte) (Usage, error) {
if status != 200 {
return Usage{}, fmt.Errorf("anthropic status %d: %w", status, ErrNotLLMResponse)
}
if isEventStream(contentType) {
return Usage{}, ErrStreamingUnsupported
}
if !isJSON(contentType) {
return Usage{}, fmt.Errorf("anthropic content-type %q: %w", contentType, ErrNotLLMResponse)
}
var resp anthropicResponse
if err := json.Unmarshal(body, &resp); err != nil {
return Usage{}, fmt.Errorf("decode anthropic response: %w: %v", ErrMalformedResponse, err)
}
return Usage{
InputTokens: resp.Usage.InputTokens,
OutputTokens: resp.Usage.OutputTokens,
TotalTokens: resp.Usage.InputTokens + resp.Usage.OutputTokens + resp.Usage.CacheReadInputTokens + resp.Usage.CacheCreationInputTokens,
CachedInputTokens: resp.Usage.CacheReadInputTokens,
CacheCreationTokens: resp.Usage.CacheCreationInputTokens,
}, nil
}
// ExtractPrompt returns the user-visible prompt text from an Anthropic
// request body. Handles the Messages API (system + messages[]) and the
// legacy /v1/complete prompt string. Returns "" on any decode failure.
func (AnthropicParser) ExtractPrompt(body []byte) string {
var req anthropicRequest
if err := json.Unmarshal(body, &req); err != nil {
return ""
}
var b strings.Builder
if len(req.System) > 0 {
if s := decodeStringOrJoin(req.System); s != "" {
b.WriteString("system: ")
b.WriteString(s)
}
}
for _, m := range req.Messages {
if b.Len() > 0 {
b.WriteByte('\n')
}
if m.Role != "" {
b.WriteString(m.Role)
b.WriteString(": ")
}
b.WriteString(decodeStringOrJoin(m.Content))
}
if b.Len() == 0 && req.Prompt != "" {
b.WriteString(req.Prompt)
}
return b.String()
}
// ExtractSessionID is the body-side fallback for Anthropic. Claude Code's
// authoritative session marker is the X-Claude-Code-Session-Id request
// header (handled by the request-parser middleware); this only mines the
// optional metadata.user_id for an embedded "...session_<uuid>" marker.
// metadata.user_id on its own is a USER identifier, not a session, so the
// whole value is deliberately NOT used — returning it would mislabel every
// request from a user as one session. Returns "" when no session marker is
// present.
func (AnthropicParser) ExtractSessionID(body []byte) string {
var req struct {
Metadata struct {
UserID string `json:"user_id"`
} `json:"metadata"`
}
if err := json.Unmarshal(body, &req); err != nil {
return ""
}
if idx := strings.LastIndex(req.Metadata.UserID, "session_"); idx >= 0 {
if session := req.Metadata.UserID[idx+len("session_"):]; session != "" {
return session
}
}
return ""
}
type anthropicMessageResponse struct {
Content []struct {
Type string `json:"type"`
Text string `json:"text"`
} `json:"content"`
// Legacy /v1/complete response.
Completion string `json:"completion"`
}
// ExtractCompletion returns the assistant text from a non-streaming Anthropic
// Messages or Completions response. Returns "" when status/content-type
// indicate the body is not parseable or no text part is present.
func (AnthropicParser) ExtractCompletion(status int, contentType string, body []byte) string {
if status != 200 || isEventStream(contentType) || !isJSON(contentType) {
return ""
}
var resp anthropicMessageResponse
if err := json.Unmarshal(body, &resp); err != nil {
return ""
}
var b strings.Builder
for _, part := range resp.Content {
if part.Text == "" {
continue
}
if b.Len() > 0 {
b.WriteByte('\n')
}
b.WriteString(part.Text)
}
if b.Len() == 0 {
return resp.Completion
}
return b.String()
}
proxy/internal/llm/anthropic_test.go
+169
View File
@@ -0,0 +1,169 @@
package llm
import (
"errors"
"os"
"path/filepath"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestAnthropicDetectFromURL(t *testing.T) {
p := AnthropicParser{}
cases := map[string]bool{
"/v1/messages": true,
"/v1/complete": true,
"/V1/Messages": true,
"/proxy/v1/messages?x": true,
"/v1/chat/completions": false,
"": false,
}
for path, want := range cases {
assert.Equal(t, want, p.DetectFromURL(path), "DetectFromURL(%q)", path)
}
}
func TestAnthropicParseRequest(t *testing.T) {
p := AnthropicParser{}
t.Run("stream true", func(t *testing.T) {
facts, err := p.ParseRequest([]byte(`{"model":"claude-sonnet-4-5","stream":true}`))
require.NoError(t, err)
assert.Equal(t, "claude-sonnet-4-5", facts.Model, "model extracted")
assert.True(t, facts.Stream, "stream flag honoured")
})
t.Run("stream default", func(t *testing.T) {
facts, err := p.ParseRequest([]byte(`{"model":"claude-sonnet-4-5"}`))
require.NoError(t, err)
assert.False(t, facts.Stream, "missing stream flag defaults to false")
})
t.Run("malformed", func(t *testing.T) {
_, err := p.ParseRequest([]byte(`{"model":`))
require.Error(t, err)
assert.True(t, errors.Is(err, ErrMalformedRequest), "sentinel wrapped")
})
}
func TestAnthropicParseResponse(t *testing.T) {
p := AnthropicParser{}
t.Run("happy fixture", func(t *testing.T) {
body, err := os.ReadFile(filepath.Join("fixtures", "anthropic_messages.json"))
require.NoError(t, err, "fixture must be readable")
usage, err := p.ParseResponse(200, "application/json", body)
require.NoError(t, err)
assert.Equal(t, int64(123), usage.InputTokens, "input tokens extracted")
assert.Equal(t, int64(45), usage.OutputTokens, "output tokens extracted")
assert.Equal(t, int64(168), usage.TotalTokens, "total computed as sum")
})
t.Run("streaming rejected", func(t *testing.T) {
_, err := p.ParseResponse(200, "text/event-stream", []byte(""))
require.ErrorIs(t, err, ErrStreamingUnsupported, "SSE responses must use the scanner")
})
t.Run("non-200", func(t *testing.T) {
_, err := p.ParseResponse(429, "application/json", []byte(`{}`))
require.ErrorIs(t, err, ErrNotLLMResponse, "non-200 rejected as non-LLM")
})
t.Run("non-json content type", func(t *testing.T) {
_, err := p.ParseResponse(200, "text/html", []byte(`{}`))
require.ErrorIs(t, err, ErrNotLLMResponse, "text/html treated as non-LLM")
})
t.Run("malformed body", func(t *testing.T) {
_, err := p.ParseResponse(200, "application/json", []byte(`{`))
require.ErrorIs(t, err, ErrMalformedResponse, "bad JSON yields malformed error")
})
// Anthropic's two cache fields are ADDITIVE to input_tokens (not
// subset). The parser must surface them so the cost meter can
// bill each bucket at its own configured rate. Total includes
// every bucket so downstream attribution sees the full token
// volume the request consumed.
t.Run("cache_read_input_tokens surfaces as CachedInputTokens (additive)", func(t *testing.T) {
body := []byte(`{"usage":{"input_tokens":256,"output_tokens":200,"cache_read_input_tokens":768}}`)
usage, err := p.ParseResponse(200, "application/json", body)
require.NoError(t, err)
assert.Equal(t, int64(256), usage.InputTokens, "regular input remains separate from cache buckets")
assert.Equal(t, int64(768), usage.CachedInputTokens, "cache_read maps onto CachedInputTokens — same field carries OpenAI cached subset and Anthropic cache reads")
assert.Zero(t, usage.CacheCreationTokens)
assert.Equal(t, int64(256+200+768), usage.TotalTokens, "total includes every input bucket plus output — cache reads are billable tokens")
})
t.Run("cache_creation_input_tokens surfaces as CacheCreationTokens (additive)", func(t *testing.T) {
body := []byte(`{"usage":{"input_tokens":256,"output_tokens":200,"cache_creation_input_tokens":512}}`)
usage, err := p.ParseResponse(200, "application/json", body)
require.NoError(t, err)
assert.Equal(t, int64(256), usage.InputTokens)
assert.Zero(t, usage.CachedInputTokens)
assert.Equal(t, int64(512), usage.CacheCreationTokens, "cache_creation surfaces — meter applies the write-rate multiplier")
assert.Equal(t, int64(256+200+512), usage.TotalTokens)
})
t.Run("both cache buckets present", func(t *testing.T) {
body := []byte(`{"usage":{"input_tokens":256,"output_tokens":200,"cache_read_input_tokens":768,"cache_creation_input_tokens":512}}`)
usage, err := p.ParseResponse(200, "application/json", body)
require.NoError(t, err)
assert.Equal(t, int64(768), usage.CachedInputTokens)
assert.Equal(t, int64(512), usage.CacheCreationTokens)
assert.Equal(t, int64(256+200+768+512), usage.TotalTokens, "all four buckets sum into total")
})
t.Run("absent cache fields leave counts at zero", func(t *testing.T) {
body := []byte(`{"usage":{"input_tokens":100,"output_tokens":50}}`)
usage, err := p.ParseResponse(200, "application/json", body)
require.NoError(t, err)
assert.Zero(t, usage.CachedInputTokens, "no cache_read field = no cached count")
assert.Zero(t, usage.CacheCreationTokens, "no cache_creation field = no creation count")
assert.Equal(t, int64(150), usage.TotalTokens, "back to the simple in+out total when no cache buckets present")
})
}
func TestAnthropicExtractPrompt_Messages(t *testing.T) {
body := []byte(`{"model":"claude-sonnet-4-7","system":"be brief","messages":[{"role":"user","content":"hi"},{"role":"assistant","content":"yes"}]}`)
got := AnthropicParser{}.ExtractPrompt(body)
require.Contains(t, got, "system: be brief", "system surfaces with role label")
require.Contains(t, got, "user: hi", "user message surfaces")
require.Contains(t, got, "assistant: yes", "assistant message surfaces")
}
func TestAnthropicExtractPrompt_LegacyComplete(t *testing.T) {
body := []byte(`{"model":"claude-2","prompt":"\n\nHuman: hi\n\nAssistant:"}`)
got := AnthropicParser{}.ExtractPrompt(body)
require.Contains(t, got, "Human: hi", "legacy prompt string surfaces")
}
func TestAnthropicExtractSessionID(t *testing.T) {
t.Run("claude code session suffix", func(t *testing.T) {
body := []byte(`{"model":"claude-opus-4-8","metadata":{"user_id":"user_abc123_account_def456_session_9f8e7d6c"},"messages":[]}`)
assert.Equal(t, "9f8e7d6c", AnthropicParser{}.ExtractSessionID(body), "session_<id> suffix must be extracted from metadata.user_id")
})
t.Run("plain user_id is not treated as a session", func(t *testing.T) {
body := []byte(`{"model":"claude-opus-4-8","metadata":{"user_id":"acme-team"},"messages":[]}`)
assert.Equal(t, "", AnthropicParser{}.ExtractSessionID(body), "a user identifier without a session marker must NOT be used as a session id")
})
t.Run("no metadata yields empty", func(t *testing.T) {
body := []byte(`{"model":"claude-opus-4-8","messages":[{"role":"user","content":"hi"}]}`)
assert.Equal(t, "", AnthropicParser{}.ExtractSessionID(body), "absent metadata.user_id yields no session id")
})
}
func TestAnthropicExtractCompletion_Messages(t *testing.T) {
body, err := os.ReadFile(filepath.Join("fixtures", "anthropic_messages.json"))
require.NoError(t, err)
got := AnthropicParser{}.ExtractCompletion(200, "application/json", body)
require.NotEmpty(t, got, "anthropic fixture has assistant text")
}
func TestAnthropicExtractCompletion_Streaming(t *testing.T) {
got := AnthropicParser{}.ExtractCompletion(200, "text/event-stream", []byte(""))
require.Empty(t, got, "streaming responses are skipped")
}
proxy/internal/llm/bedrock.go
+193
View File
@@ -0,0 +1,193 @@
package llm
import (
"encoding/json"
"fmt"
"strings"
)
// ProviderNameBedrock is the stable label for the AWS Bedrock parser, used as
// the llm.provider metadata value and the cost-meter formula selector.
const ProviderNameBedrock = "bedrock"
// BedrockParser implements the Parser interface for the AWS Bedrock runtime.
// Bedrock carries the model in the URL path (/model/{id}/{action}); the request
// middleware extracts it there, so this parser focuses on the response shapes:
// the vendor-native InvokeModel body (e.g. Anthropic's snake_case usage) and the
// unified Converse body (camelCase usage).
type BedrockParser struct{}
var bedrockPathHints = []string{"/invoke", "/converse"}
// Provider returns ProviderBedrock.
func (BedrockParser) Provider() Provider { return ProviderBedrock }
// ProviderName returns the stable label used for metrics and metadata.
func (BedrockParser) ProviderName() string { return ProviderNameBedrock }
// DetectFromURL reports whether the path is a Bedrock runtime model endpoint.
func (BedrockParser) DetectFromURL(path string) bool {
lower := strings.ToLower(path)
if !strings.HasPrefix(lower, "/model/") {
return false
}
for _, hint := range bedrockPathHints {
if strings.Contains(lower, hint) {
return true
}
}
return false
}
// ParseRequest is a no-op for Bedrock: the model lives in the URL path, not the
// body, and the streaming flag is derived from the path action. The request
// middleware handles both via parseBedrockPath, so this returns empty facts.
func (BedrockParser) ParseRequest([]byte) (RequestFacts, error) {
return RequestFacts{}, nil
}
// bedrockResponse captures token usage from both Bedrock response shapes:
// InvokeModel (vendor-native; Anthropic uses snake_case + additive cache
// buckets) and Converse (camelCase, with a precomputed total).
type bedrockResponse struct {
Usage struct {
// InvokeModel (Anthropic-on-Bedrock) — snake_case.
InputTokens int64 `json:"input_tokens"`
OutputTokens int64 `json:"output_tokens"`
CacheReadInputTokens int64 `json:"cache_read_input_tokens"`
CacheCreationInputTokens int64 `json:"cache_creation_input_tokens"`
// Converse — camelCase; cache buckets are additive to inputTokens (AWS names the write bucket cacheWriteInputTokens).
InputTokensCamel int64 `json:"inputTokens"`
OutputTokensCamel int64 `json:"outputTokens"`
TotalTokensCamel int64 `json:"totalTokens"`
CacheReadTokensCamel int64 `json:"cacheReadInputTokens"`
CacheWriteTokensCamel int64 `json:"cacheWriteInputTokens"`
} `json:"usage"`
}
// ParseResponse decodes the non-streaming Bedrock response envelope, handling
// both the InvokeModel and Converse usage shapes. Non-200 / non-JSON bodies are
// treated as non-LLM responses so the caller skips cost accounting.
func (BedrockParser) ParseResponse(status int, contentType string, body []byte) (Usage, error) {
if status != 200 {
return Usage{}, fmt.Errorf("bedrock status %d: %w", status, ErrNotLLMResponse)
}
if isAWSEventStream(contentType) || isEventStream(contentType) {
return Usage{}, ErrStreamingUnsupported
}
if !isJSON(contentType) {
return Usage{}, fmt.Errorf("bedrock content-type %q: %w", contentType, ErrNotLLMResponse)
}
var resp bedrockResponse
if err := json.Unmarshal(body, &resp); err != nil {
return Usage{}, fmt.Errorf("decode bedrock response: %w: %v", ErrMalformedResponse, err)
}
inTok := firstNonZero(resp.Usage.InputTokens, resp.Usage.InputTokensCamel)
outTok := firstNonZero(resp.Usage.OutputTokens, resp.Usage.OutputTokensCamel)
cacheRead := firstNonZero(resp.Usage.CacheReadInputTokens, resp.Usage.CacheReadTokensCamel)
cacheWrite := firstNonZero(resp.Usage.CacheCreationInputTokens, resp.Usage.CacheWriteTokensCamel)
total := resp.Usage.TotalTokensCamel
if total == 0 {
total = inTok + outTok + cacheRead + cacheWrite
}
return Usage{
InputTokens: inTok,
OutputTokens: outTok,
TotalTokens: total,
CachedInputTokens: cacheRead,
CacheCreationTokens: cacheWrite,
}, nil
}
// ExtractPrompt returns the user-visible prompt from a Bedrock request body,
// handling both the InvokeModel (Anthropic Messages: system + messages[]) and
// Converse (messages[].content[].text) shapes. Returns "" on decode failure.
func (BedrockParser) ExtractPrompt(body []byte) string {
var req struct {
System json.RawMessage `json:"system"`
Messages []struct {
Role string `json:"role"`
Content json.RawMessage `json:"content"`
} `json:"messages"`
}
if err := json.Unmarshal(body, &req); err != nil {
return ""
}
var b strings.Builder
if s := decodeStringOrJoin(req.System); s != "" {
b.WriteString("system: ")
b.WriteString(s)
}
for _, m := range req.Messages {
if b.Len() > 0 {
b.WriteByte('\n')
}
if m.Role != "" {
b.WriteString(m.Role)
b.WriteString(": ")
}
b.WriteString(decodeStringOrJoin(m.Content))
}
return b.String()
}
// ExtractCompletion returns the assistant text from a non-streaming Bedrock
// response, handling InvokeModel (Anthropic content[].text) and Converse
// (output.message.content[].text).
func (BedrockParser) ExtractCompletion(status int, contentType string, body []byte) string {
if status != 200 || isAWSEventStream(contentType) || !isJSON(contentType) {
return ""
}
var resp struct {
Content []struct {
Text string `json:"text"`
} `json:"content"`
Output struct {
Message struct {
Content []struct {
Text string `json:"text"`
} `json:"content"`
} `json:"message"`
} `json:"output"`
}
if err := json.Unmarshal(body, &resp); err != nil {
return ""
}
var b strings.Builder
appendText := func(text string) {
if text == "" {
return
}
if b.Len() > 0 {
b.WriteByte('\n')
}
b.WriteString(text)
}
for _, p := range resp.Content {
appendText(p.Text)
}
for _, p := range resp.Output.Message.Content {
appendText(p.Text)
}
return b.String()
}
// ExtractSessionID has no Bedrock-native marker; session grouping relies on the
// request headers handled by the middleware. Returns "".
func (BedrockParser) ExtractSessionID([]byte) string { return "" }
// firstNonZero returns a when non-zero, else b. Folds the snake_case and
// camelCase usage variants into a single value.
func firstNonZero(a, b int64) int64 {
if a != 0 {
return a
}
return b
}
// isAWSEventStream reports whether contentType is the AWS binary event-stream
// framing used by Bedrock's streaming endpoints.
func isAWSEventStream(contentType string) bool {
return strings.Contains(strings.ToLower(contentType), "application/vnd.amazon.eventstream")
}
proxy/internal/llm/bedrock_model.go
+38
View File
@@ -0,0 +1,38 @@
package llm
import (
"regexp"
"strings"
)
// bedrockRegionPrefixes are the cross-region inference-profile prefixes that
// front a Bedrock model id (e.g. "eu.anthropic.claude-...").
var bedrockRegionPrefixes = []string{"us.", "eu.", "apac.", "global."}
// bedrockVersionSuffix matches the trailing "-vN[:N]" or "-YYYYMMDD-vN[:N]"
// version/throughput suffix of a Bedrock model id.
var bedrockVersionSuffix = regexp.MustCompile(`-(\d{8}-)?v\d+(:\d+)?$`)
// NormalizeBedrockModel strips an ARN wrapper, a cross-region inference-profile
// prefix, and the version/throughput suffix from a Bedrock model id so it
// matches the catalog/pricing key, e.g.
// "eu.anthropic.claude-sonnet-4-5-20250929-v1:0" -> "anthropic.claude-sonnet-4-5"
// and the inference-profile ARN's last segment likewise. It is the single
// source of truth shared by the request parser (which normalizes the request
// model from the URL path) and the router (which normalizes the operator's
// registered Bedrock model ids so both sides compare equal).
func NormalizeBedrockModel(modelID string) string {
m := modelID
if strings.HasPrefix(m, "arn:") {
if i := strings.LastIndex(m, "/"); i >= 0 {
m = m[i+1:]
}
}
for _, p := range bedrockRegionPrefixes {
if strings.HasPrefix(m, p) {
m = m[len(p):]
break
}
}
return bedrockVersionSuffix.ReplaceAllString(m, "")
}
proxy/internal/llm/bedrock_model_test.go
+23
View File
@@ -0,0 +1,23 @@
package llm
import (
"testing"
"github.com/stretchr/testify/require"
)
func TestNormalizeBedrockModel(t *testing.T) {
cases := map[string]string{
"eu.anthropic.claude-sonnet-4-5-20250929-v1:0": "anthropic.claude-sonnet-4-5",
"us.anthropic.claude-haiku-4-5": "anthropic.claude-haiku-4-5",
"us.anthropic.claude-opus-4-8-20250101-v1:0": "anthropic.claude-opus-4-8",
"anthropic.claude-sonnet-4-5-20250929-v1:0": "anthropic.claude-sonnet-4-5",
"meta.llama3-3-70b-instruct-v1:0": "meta.llama3-3-70b-instruct",
"amazon.nova-pro-v1:0": "amazon.nova-pro",
// Inference-profile ARN — model id lives in the last path segment.
"arn:aws:bedrock:eu-central-1:123456789012:inference-profile/eu.anthropic.claude-sonnet-4-5-20250929-v1:0": "anthropic.claude-sonnet-4-5",
}
for in, want := range cases {
require.Equal(t, want, NormalizeBedrockModel(in), "normalize %q", in)
}
}
proxy/internal/llm/bedrock_test.go
+77
View File
@@ -0,0 +1,77 @@
package llm
import (
"testing"
"github.com/stretchr/testify/require"
)
func TestBedrockParser_ParseResponse_Invoke(t *testing.T) {
body := []byte(`{"usage":{"input_tokens":13,"output_tokens":5,"cache_read_input_tokens":2,"cache_creation_input_tokens":4}}`)
u, err := BedrockParser{}.ParseResponse(200, "application/json", body)
require.NoError(t, err)
require.Equal(t, int64(13), u.InputTokens, "invoke input tokens")
require.Equal(t, int64(5), u.OutputTokens, "invoke output tokens")
require.Equal(t, int64(2), u.CachedInputTokens, "invoke cache-read tokens")
require.Equal(t, int64(4), u.CacheCreationTokens, "invoke cache-creation tokens")
require.Equal(t, int64(13+5+2+4), u.TotalTokens, "invoke total is additive")
}
func TestBedrockParser_ParseResponse_Converse(t *testing.T) {
body := []byte(`{"output":{"message":{"content":[{"text":"pong"}]}},"usage":{"inputTokens":11,"outputTokens":3,"totalTokens":14}}`)
u, err := BedrockParser{}.ParseResponse(200, "application/json", body)
require.NoError(t, err)
require.Equal(t, int64(11), u.InputTokens, "converse camelCase input tokens")
require.Equal(t, int64(3), u.OutputTokens, "converse camelCase output tokens")
require.Equal(t, int64(14), u.TotalTokens, "converse uses provider total")
}
// Converse camelCase cache fields must land in the billed Usage buckets, same as the InvokeModel snake_case fields.
func TestBedrockParser_ParseResponse_ConverseCacheBuckets(t *testing.T) {
body := []byte(`{"usage":{"inputTokens":11,"outputTokens":3,"cacheReadInputTokens":7,"cacheWriteInputTokens":9}}`)
u, err := BedrockParser{}.ParseResponse(200, "application/json", body)
require.NoError(t, err)
require.Equal(t, int64(11), u.InputTokens, "converse input tokens")
require.Equal(t, int64(3), u.OutputTokens, "converse output tokens")
require.Equal(t, int64(7), u.CachedInputTokens, "converse cache-read tokens")
require.Equal(t, int64(9), u.CacheCreationTokens, "converse cache-write tokens")
require.Equal(t, int64(11+3+7+9), u.TotalTokens, "total backfill is additive when the provider omits totalTokens")
}
func TestBedrockParser_ParseResponse_StreamingUnsupported(t *testing.T) {
_, err := BedrockParser{}.ParseResponse(200, "application/vnd.amazon.eventstream", []byte("binary"))
require.ErrorIs(t, err, ErrStreamingUnsupported, "event-stream must route to the streaming accumulator")
}
func TestBedrockParser_ParseResponse_NonSuccess(t *testing.T) {
_, err := BedrockParser{}.ParseResponse(404, "application/json", []byte(`{"message":"gated"}`))
require.ErrorIs(t, err, ErrNotLLMResponse, "non-200 is not an LLM response")
}
func TestBedrockParser_ExtractCompletion(t *testing.T) {
invoke := BedrockParser{}.ExtractCompletion(200, "application/json", []byte(`{"content":[{"text":"a"},{"text":"b"}]}`))
require.Equal(t, "a\nb", invoke, "invoke completion joins content parts")
converse := BedrockParser{}.ExtractCompletion(200, "application/json", []byte(`{"output":{"message":{"content":[{"text":"x"}]}}}`))
require.Equal(t, "x", converse, "converse completion reads output.message.content")
}
func TestBedrockParser_ExtractPrompt(t *testing.T) {
invoke := BedrockParser{}.ExtractPrompt([]byte(`{"messages":[{"role":"user","content":"hi"}]}`))
require.Equal(t, "user: hi", invoke, "invoke prompt reads anthropic content string")
converse := BedrockParser{}.ExtractPrompt([]byte(`{"messages":[{"role":"user","content":[{"text":"hello"}]}]}`))
require.Equal(t, "user: hello", converse, "converse prompt reads content parts")
}
func TestBedrockParser_DetectFromURL(t *testing.T) {
require.True(t, BedrockParser{}.DetectFromURL("/model/eu.anthropic.claude/invoke"), "invoke path")
require.True(t, BedrockParser{}.DetectFromURL("/model/x/converse-stream"), "converse-stream path")
require.False(t, BedrockParser{}.DetectFromURL("/v1/chat/completions"), "openai path is not bedrock")
}
func TestBedrockParser_RegisteredByName(t *testing.T) {
p, ok := ParserByName(ProviderNameBedrock)
require.True(t, ok, "bedrock parser is registered")
require.Equal(t, ProviderNameBedrock, p.ProviderName())
}
proxy/internal/llm/errors.go
+31
View File
@@ -0,0 +1,31 @@
package llm
import "errors"
// Sentinel errors returned by parsers and the pricing loader. Callers use
// errors.Is to branch on a condition without coupling to parser internals.
var (
// ErrUnknownProvider indicates no parser claimed the request path.
ErrUnknownProvider = errors.New("llmobs: unknown provider")
// ErrUnsupportedModel indicates the response parsed successfully but the
// model is absent from the pricing table. Token counts are still valid.
ErrUnsupportedModel = errors.New("llmobs: unsupported model")
// ErrNotLLMResponse indicates the response is not a JSON success body
// that a non-streaming parser can consume (non-200 or wrong content type).
ErrNotLLMResponse = errors.New("llmobs: not an LLM response")
// ErrStreamingUnsupported indicates the caller passed an SSE response to
// a non-streaming parser. Streaming is handled separately via the SSE
// scanner.
ErrStreamingUnsupported = errors.New("llmobs: streaming response requires SSE scanner")
// ErrMalformedResponse indicates the response body could not be decoded
// as the provider-specific JSON schema.
ErrMalformedResponse = errors.New("llmobs: malformed response body")
// ErrMalformedRequest indicates the request body could not be decoded as
// the provider-specific JSON schema.
ErrMalformedRequest = errors.New("llmobs: malformed request body")
)
proxy/internal/llm/fixtures/anthropic_messages.json
+17
View File
@@ -0,0 +1,17 @@
{
"id": "msg_abc",
"type": "message",
"role": "assistant",
"model": "claude-sonnet-4-5",
"content": [
{
"type": "text",
"text": "Hello, world!"
}
],
"stop_reason": "end_turn",
"usage": {
"input_tokens": 123,
"output_tokens": 45
}
}
proxy/internal/llm/fixtures/anthropic_stream.txt
+21
View File
@@ -0,0 +1,21 @@
event: message_start
data: {"type":"message_start","message":{"id":"msg_abc","type":"message","role":"assistant","model":"claude-sonnet-4-5","content":[],"stop_reason":null,"usage":{"input_tokens":123,"output_tokens":1}}}
event: content_block_start
data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}
event: content_block_delta
data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hello"}}
event: content_block_delta
data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":", world!"}}
event: content_block_stop
data: {"type":"content_block_stop","index":0}
event: message_delta
data: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":45}}
event: message_stop
data: {"type":"message_stop"}
proxy/internal/llm/fixtures/openai_chat_completion.json
+21
View File
@@ -0,0 +1,21 @@
{
"id": "chatcmpl-abc",
"object": "chat.completion",
"created": 1700000000,
"model": "gpt-4o-mini",
"choices": [
{
"index": 0,
"message": {
"role": "assistant",
"content": "Hello, world!"
},
"finish_reason": "stop"
}
],
"usage": {
"prompt_tokens": 123,
"completion_tokens": 45,
"total_tokens": 168
}
}
proxy/internal/llm/fixtures/openai_responses.json
+24
View File
@@ -0,0 +1,24 @@
{
"id": "resp_abc",
"object": "response",
"created_at": 1700000000,
"model": "gpt-5.4",
"output": [
{
"type": "message",
"role": "assistant",
"content": [{"type": "output_text", "text": "ok"}]
}
],
"usage": {
"input_tokens": 15,
"input_tokens_details": {
"cached_tokens": 0
},
"output_tokens": 414,
"output_tokens_details": {
"reasoning_tokens": 0
},
"total_tokens": 429
}
}
proxy/internal/llm/fixtures/openai_responses_stream.txt
+24
View File
@@ -0,0 +1,24 @@
event: response.created
data: {"type":"response.created","response":{"id":"resp_abc","object":"response","model":"gpt-5.5","usage":null}}
event: response.in_progress
data: {"type":"response.in_progress","response":{"id":"resp_abc","usage":null}}
event: response.output_item.added
data: {"type":"response.output_item.added","output_index":0,"item":{"type":"message","role":"assistant","content":[]}}
event: response.content_part.added
data: {"type":"response.content_part.added","item_id":"msg_1","output_index":0,"content_index":0,"part":{"type":"output_text","text":""}}
event: response.output_text.delta
data: {"type":"response.output_text.delta","item_id":"msg_1","output_index":0,"content_index":0,"delta":"Hello"}
event: response.output_text.delta
data: {"type":"response.output_text.delta","item_id":"msg_1","output_index":0,"content_index":0,"delta":", world!"}
event: response.output_text.done
data: {"type":"response.output_text.done","item_id":"msg_1","output_index":0,"content_index":0,"text":"Hello, world!"}
event: response.completed
data: {"type":"response.completed","response":{"id":"resp_abc","object":"response","model":"gpt-5.5","usage":{"input_tokens":123,"input_tokens_details":{"cached_tokens":40},"output_tokens":45,"output_tokens_details":{"reasoning_tokens":12},"total_tokens":168}}}
proxy/internal/llm/fixtures/openai_stream.txt
+8
View File
@@ -0,0 +1,8 @@
data: {"id":"chatcmpl-abc","object":"chat.completion.chunk","created":1700000000,"model":"gpt-4o-mini","choices":[{"index":0,"delta":{"role":"assistant","content":"Hello"},"finish_reason":null}]}
data: {"id":"chatcmpl-abc","object":"chat.completion.chunk","created":1700000000,"model":"gpt-4o-mini","choices":[{"index":0,"delta":{"content":", world!"},"finish_reason":null}]}
data: {"id":"chatcmpl-abc","object":"chat.completion.chunk","created":1700000000,"model":"gpt-4o-mini","choices":[{"index":0,"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":123,"completion_tokens":45,"total_tokens":168}}
data: [DONE]
proxy/internal/llm/fixtures/pricing.yaml
+59
View File
@@ -0,0 +1,59 @@
# Realistic-pricing starter for llm_observability. Drop this into the
# directory you point the proxy at via --plugin-data-dir, then reference it
# from the target's plugin config:
#
# plugins:
# - id: llm_observability
# enabled: true
# params:
# pricing_path: pricing.yaml
#
# Values are USD per 1_000 tokens. Public list prices drift; treat this as a
# starting point and keep your production copy current.
openai:
# GPT-5 family
gpt-5:
input_per_1k: 0.00125
output_per_1k: 0.01
gpt-5-mini:
input_per_1k: 0.00025
output_per_1k: 0.002
gpt-5-nano:
input_per_1k: 0.00005
output_per_1k: 0.0004
gpt-5.4:
input_per_1k: 0.00125
output_per_1k: 0.01
# GPT-4o family
gpt-4o:
input_per_1k: 0.0025
output_per_1k: 0.01
gpt-4o-mini:
input_per_1k: 0.00015
output_per_1k: 0.0006
# Embeddings
text-embedding-3-large:
input_per_1k: 0.00013
output_per_1k: 0
text-embedding-3-small:
input_per_1k: 0.00002
output_per_1k: 0
anthropic:
# Claude 4.x family
claude-opus-4-7:
input_per_1k: 0.015
output_per_1k: 0.075
claude-sonnet-4-7:
input_per_1k: 0.003
output_per_1k: 0.015
claude-sonnet-4-6:
input_per_1k: 0.003
output_per_1k: 0.015
claude-sonnet-4-5:
input_per_1k: 0.003
output_per_1k: 0.015
claude-haiku-4-5:
input_per_1k: 0.0008
output_per_1k: 0.004
proxy/internal/llm/openai.go
+412
View File
@@ -0,0 +1,412 @@
package llm
import (
"encoding/json"
"fmt"
"strings"
)
// OpenAIParser implements the Parser interface for OpenAI-compatible APIs.
// It recognizes chat.completions, completions, embeddings, and the newer
// responses endpoint; any proxy path-prefix stripping is tolerated by the
// substring match in DetectFromURL.
type OpenAIParser struct{}
// openAIPathHints are substring patterns that mark a request as
// OpenAI-shaped. The bare `/chat/completions` is listed alongside
// `/v1/chat/completions` because gateways like Cloudflare AI
// Gateway place their own version segment before the provider
// slug (gateway/v1/{account}/{gateway}/openai/chat/completions) —
// the canonical `/v1/` ends up nowhere near `/chat/completions`,
// so the `/v1/chat/completions` hint misses. `/chat/completions`
// is OpenAI's API contract: any service accepting OpenAI bodies
// serves at this path, so false-positive risk is negligible.
// `/completions` (legacy), `/embeddings`, and `/responses` are
// kept on the canonical-only path because their bare forms are
// too generic to be safe substrings.
var openAIPathHints = []string{
"/v1/chat/completions",
"/v1/completions",
"/v1/embeddings",
"/v1/responses",
"/chat/completions",
}
// Provider returns ProviderOpenAI.
func (OpenAIParser) Provider() Provider { return ProviderOpenAI }
// ProviderName returns the stable label used for metrics and metadata.
func (OpenAIParser) ProviderName() string { return "openai" }
// DetectFromURL reports whether the given request path looks like an OpenAI
// API endpoint. The match is case-insensitive and substring-based so that a
// reverse proxy prefix strip or rewrite does not defeat detection.
func (OpenAIParser) DetectFromURL(path string) bool {
lower := strings.ToLower(path)
for _, hint := range openAIPathHints {
if strings.Contains(lower, hint) {
return true
}
}
return false
}
type openAIRequest struct {
Model string `json:"model"`
Stream *bool `json:"stream"`
StreamOptions *struct {
IncludeUsage *bool `json:"include_usage"`
} `json:"stream_options"`
// Chat Completions / Completions: messages[].content (string or array of
// content parts). Responses API: input is either a string or an array of
// items with content parts. We use json.RawMessage to defer parsing each
// shape independently.
Messages []openAIMessage `json:"messages"`
Prompt json.RawMessage `json:"prompt"`
Input json.RawMessage `json:"input"`
}
type openAIMessage struct {
Role string `json:"role"`
Content json.RawMessage `json:"content"`
}
// ParseRequest extracts the model name and streaming flag from an OpenAI
// request body. Unknown or missing fields leave the corresponding struct
// members zero-valued.
func (OpenAIParser) ParseRequest(body []byte) (RequestFacts, error) {
var req openAIRequest
if err := json.Unmarshal(body, &req); err != nil {
return RequestFacts{}, fmt.Errorf("decode openai request: %w: %v", ErrMalformedRequest, err)
}
return RequestFacts{
Model: req.Model,
Stream: ptrDeref(req.Stream),
}, nil
}
// openAIResponse accepts both naming conventions in a single struct because
// OpenAI's older Chat Completions API uses prompt_tokens/completion_tokens
// while the newer Responses API (/v1/responses) uses input_tokens/output_tokens
// (aligned with Anthropic). Pointer fields let us tell "absent" from "zero".
//
// PromptTokensDetails.CachedTokens (Chat Completions) and
// InputTokensDetails.CachedTokens (Responses API) carry the SUBSET of
// prompt/input tokens that hit the prompt cache. Cost-meter applies the
// discount rate to that subset and the regular rate to the remainder so
// we never double-bill the cached portion.
type openAIResponse struct {
Usage struct {
PromptTokens *int64 `json:"prompt_tokens"`
CompletionTokens *int64 `json:"completion_tokens"`
InputTokens *int64 `json:"input_tokens"`
OutputTokens *int64 `json:"output_tokens"`
TotalTokens *int64 `json:"total_tokens"`
PromptTokensDetails *struct {
CachedTokens *int64 `json:"cached_tokens"`
} `json:"prompt_tokens_details"`
InputTokensDetails *struct {
CachedTokens *int64 `json:"cached_tokens"`
} `json:"input_tokens_details"`
} `json:"usage"`
}
// ParseResponse decodes the non-streaming OpenAI response envelope. Status
// codes other than 200 are treated as non-LLM responses so the caller can
// skip cost accounting without aborting the request.
func (OpenAIParser) ParseResponse(status int, contentType string, body []byte) (Usage, error) {
if status != 200 {
return Usage{}, fmt.Errorf("openai status %d: %w", status, ErrNotLLMResponse)
}
if isEventStream(contentType) {
return Usage{}, ErrStreamingUnsupported
}
if !isJSON(contentType) {
return Usage{}, fmt.Errorf("openai content-type %q: %w", contentType, ErrNotLLMResponse)
}
var resp openAIResponse
if err := json.Unmarshal(body, &resp); err != nil {
return Usage{}, fmt.Errorf("decode openai response: %w: %v", ErrMalformedResponse, err)
}
// Responses-API names take precedence when present; fall back to the older
// Chat Completions names. This handles both endpoints transparently
// without forcing a per-route configuration.
u := Usage{
InputTokens: pickInt64(resp.Usage.InputTokens, resp.Usage.PromptTokens),
OutputTokens: pickInt64(resp.Usage.OutputTokens, resp.Usage.CompletionTokens),
TotalTokens: derefInt64(resp.Usage.TotalTokens),
CachedInputTokens: openAICachedTokens(resp),
}
if u.TotalTokens == 0 && (u.InputTokens > 0 || u.OutputTokens > 0) {
u.TotalTokens = u.InputTokens + u.OutputTokens
}
return u, nil
}
// openAICachedTokens returns the cached-prompt subset reported by
// either the Responses-API (input_tokens_details.cached_tokens) or
// the Chat-Completions API (prompt_tokens_details.cached_tokens).
// Responses-API takes precedence when both are populated.
func openAICachedTokens(resp openAIResponse) int64 {
// Responses-API details are authoritative when present: an explicit
// cached_tokens of 0 must be honored, not treated as missing and
// overridden by the Chat-Completions field (which would overstate cache).
if resp.Usage.InputTokensDetails != nil && resp.Usage.InputTokensDetails.CachedTokens != nil {
return derefInt64(resp.Usage.InputTokensDetails.CachedTokens)
}
if resp.Usage.PromptTokensDetails != nil {
return derefInt64(resp.Usage.PromptTokensDetails.CachedTokens)
}
return 0
}
// ExtractPrompt returns the user-visible prompt text from an OpenAI request.
// Handles chat.completions (messages[].content), legacy completions (prompt
// string), and the Responses API (input as string or content-part array).
// Returns "" when nothing extractable is found.
func (OpenAIParser) ExtractPrompt(body []byte) string {
var req openAIRequest
if err := json.Unmarshal(body, &req); err != nil {
return ""
}
if len(req.Messages) > 0 {
return joinMessages(req.Messages)
}
if len(req.Input) > 0 {
return extractResponsesInput(req.Input)
}
if len(req.Prompt) > 0 {
return decodeStringOrJoin(req.Prompt)
}
return ""
}
// extractResponsesInput handles the Responses API `input` field. It is one
// of three shapes: a plain string, an array of message items
// ({role, content: string | [parts]}) as sent by Codex and the Responses
// SDK, or a flat array of content parts ({type, text/input_text}). Message
// items are flattened to "role: text" lines; items without extractable text
// (reasoning blocks, tool calls) are skipped.
func extractResponsesInput(raw json.RawMessage) string {
if s, ok := tryDecodeString(raw); ok {
return s
}
var items []struct {
Role string `json:"role"`
Content json.RawMessage `json:"content"`
Text string `json:"text"`
InputText string `json:"input_text"`
}
if err := json.Unmarshal(raw, &items); err != nil {
return extractContentParts(raw)
}
var b strings.Builder
for _, it := range items {
var text string
switch {
case len(it.Content) > 0:
text = decodeStringOrJoin(it.Content)
case it.Text != "":
text = it.Text
case it.InputText != "":
text = it.InputText
}
if text == "" {
continue
}
if b.Len() > 0 {
b.WriteByte('\n')
}
if it.Role != "" {
b.WriteString(it.Role)
b.WriteString(": ")
}
b.WriteString(text)
}
return b.String()
}
// ExtractSessionID reads the OpenAI session marker. Codex (the Responses
// API client) stamps client_metadata.session_id on every request body;
// plain chat.completions traffic carries no session id and yields "".
func (OpenAIParser) ExtractSessionID(body []byte) string {
var req struct {
ClientMetadata struct {
SessionID string `json:"session_id"`
} `json:"client_metadata"`
}
if err := json.Unmarshal(body, &req); err != nil {
return ""
}
return req.ClientMetadata.SessionID
}
type openAIChatChoice struct {
Message struct {
Role string `json:"role"`
Content json.RawMessage `json:"content"`
} `json:"message"`
Text string `json:"text"`
}
type openAIChatResponse struct {
Choices []openAIChatChoice `json:"choices"`
// Responses API: output[].content[].text
Output []struct {
Type string `json:"type"`
Content json.RawMessage `json:"content"`
Text string `json:"text"`
} `json:"output"`
OutputText string `json:"output_text"`
}
// ExtractCompletion returns the assistant text from a non-streaming OpenAI
// response. Handles chat.completions (choices[].message.content), legacy
// completions (choices[].text), and Responses API (output[].content[].text
// or the convenience output_text field).
func (OpenAIParser) ExtractCompletion(status int, contentType string, body []byte) string {
if status != 200 || isEventStream(contentType) || !isJSON(contentType) {
return ""
}
var resp openAIChatResponse
if err := json.Unmarshal(body, &resp); err != nil {
return ""
}
if resp.OutputText != "" {
return resp.OutputText
}
for _, c := range resp.Choices {
if len(c.Message.Content) > 0 {
if s := decodeStringOrJoin(c.Message.Content); s != "" {
return s
}
}
if c.Text != "" {
return c.Text
}
}
for _, o := range resp.Output {
if o.Text != "" {
return o.Text
}
if len(o.Content) > 0 {
if s := extractContentParts(o.Content); s != "" {
return s
}
}
}
return ""
}
// joinMessages flattens a chat.completions messages array into a single
// "role: content" string per message, separated by newlines. Roles surface
// system/user/assistant context which is useful for log review.
func joinMessages(msgs []openAIMessage) string {
var b strings.Builder
for i, m := range msgs {
if i > 0 {
b.WriteByte('\n')
}
if m.Role != "" {
b.WriteString(m.Role)
b.WriteString(": ")
}
b.WriteString(decodeStringOrJoin(m.Content))
}
return b.String()
}
// extractContentParts handles the Responses-API content shape, which is
// either a single string or an array of {type, text} parts. text and
// input_text both carry user-facing content.
func extractContentParts(raw json.RawMessage) string {
if s, ok := tryDecodeString(raw); ok {
return s
}
var parts []struct {
Type string `json:"type"`
Text string `json:"text"`
InputText string `json:"input_text"`
}
if err := json.Unmarshal(raw, &parts); err != nil {
// Last-ditch: array of strings.
var arr []string
if json.Unmarshal(raw, &arr) == nil {
return strings.Join(arr, "\n")
}
return ""
}
var b strings.Builder
for _, p := range parts {
var text string
switch {
case p.Text != "":
text = p.Text
case p.InputText != "":
text = p.InputText
}
if text == "" {
continue
}
if b.Len() > 0 {
b.WriteByte('\n')
}
b.WriteString(text)
}
return b.String()
}
// decodeStringOrJoin accepts either a JSON string or a content-parts array
// (chat.completions multimodal) and returns a flat string. Multimodal parts
// are separated by newlines; non-text parts are skipped.
func decodeStringOrJoin(raw json.RawMessage) string {
if s, ok := tryDecodeString(raw); ok {
return s
}
return extractContentParts(raw)
}
func tryDecodeString(raw json.RawMessage) (string, bool) {
if len(raw) == 0 {
return "", false
}
var s string
if err := json.Unmarshal(raw, &s); err == nil {
return s, true
}
return "", false
}
// pickInt64 returns the first non-nil pointer's value. Used to prefer one
// naming convention while transparently falling back to another.
func pickInt64(preferred, fallback *int64) int64 {
if preferred != nil {
return *preferred
}
return derefInt64(fallback)
}
func derefInt64(v *int64) int64 {
if v == nil {
return 0
}
return *v
}
func ptrDeref(b *bool) bool {
if b == nil {
return false
}
return *b
}
func isEventStream(contentType string) bool {
return strings.Contains(strings.ToLower(contentType), "text/event-stream")
}
func isJSON(contentType string) bool {
lower := strings.ToLower(contentType)
return strings.Contains(lower, "application/json") || strings.Contains(lower, "+json")
}
proxy/internal/llm/openai_test.go
+255
View File
@@ -0,0 +1,255 @@
package llm
import (
"errors"
"os"
"path/filepath"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestOpenAIDetectFromURL(t *testing.T) {
p := OpenAIParser{}
cases := map[string]bool{
"/v1/chat/completions": true,
"/v1/completions": true,
"/v1/embeddings": true,
"/v1/responses": true,
"/API/V1/Chat/Completions": true,
"/upstream/v1/chat/completions?trace=1": true,
// Cloudflare AI Gateway puts its own /v1/{account}/{gateway}
// segment between the canonical /v1/ and the provider slug,
// so the /v1/chat/completions substring no longer appears
// adjacent in the path. The bare /chat/completions hint
// catches Cloudflare's OpenAI direct path
// (/v1/{account}/{gateway}/openai/chat/completions) and
// compat path (/v1/{account}/{gateway}/compat/chat/completions).
"/v1/{account}/{gateway}/openai/chat/completions": true,
"/v1/{account}/{gateway}/compat/chat/completions": true,
"/chat/completions": true,
"/v1/messages": false,
"/healthz": false,
"": false,
}
for path, want := range cases {
assert.Equal(t, want, p.DetectFromURL(path), "DetectFromURL(%q)", path)
}
}
func TestOpenAIParseRequest(t *testing.T) {
p := OpenAIParser{}
t.Run("stream true", func(t *testing.T) {
facts, err := p.ParseRequest([]byte(`{"model":"gpt-4o","stream":true,"stream_options":{"include_usage":true}}`))
require.NoError(t, err)
assert.Equal(t, "gpt-4o", facts.Model, "request model extracted")
assert.True(t, facts.Stream, "request marked as streaming")
})
t.Run("stream default", func(t *testing.T) {
facts, err := p.ParseRequest([]byte(`{"model":"gpt-4o-mini"}`))
require.NoError(t, err)
assert.Equal(t, "gpt-4o-mini", facts.Model, "request model extracted")
assert.False(t, facts.Stream, "missing stream flag defaults to false")
})
t.Run("malformed", func(t *testing.T) {
_, err := p.ParseRequest([]byte(`{not json}`))
require.Error(t, err)
assert.True(t, errors.Is(err, ErrMalformedRequest), "sentinel error wrapped")
})
}
func TestOpenAIParseResponse(t *testing.T) {
p := OpenAIParser{}
t.Run("happy fixture", func(t *testing.T) {
body, err := os.ReadFile(filepath.Join("fixtures", "openai_chat_completion.json"))
require.NoError(t, err, "fixture must be readable")
usage, err := p.ParseResponse(200, "application/json", body)
require.NoError(t, err)
assert.Equal(t, int64(123), usage.InputTokens, "prompt tokens become input")
assert.Equal(t, int64(45), usage.OutputTokens, "completion tokens become output")
assert.Equal(t, int64(168), usage.TotalTokens, "total_tokens carried through")
})
t.Run("total computed when missing", func(t *testing.T) {
body := []byte(`{"usage":{"prompt_tokens":10,"completion_tokens":5}}`)
usage, err := p.ParseResponse(200, "application/json", body)
require.NoError(t, err)
assert.Equal(t, int64(15), usage.TotalTokens, "total computed from in+out")
})
t.Run("streaming rejected", func(t *testing.T) {
_, err := p.ParseResponse(200, "text/event-stream", []byte(""))
require.ErrorIs(t, err, ErrStreamingUnsupported, "SSE responses must use the scanner")
})
t.Run("non-200", func(t *testing.T) {
_, err := p.ParseResponse(500, "application/json", []byte(`{"error":"x"}`))
require.ErrorIs(t, err, ErrNotLLMResponse, "non-200 rejected as non-LLM")
})
t.Run("non-json content type", func(t *testing.T) {
_, err := p.ParseResponse(200, "text/plain", []byte(`{}`))
require.ErrorIs(t, err, ErrNotLLMResponse, "text/plain treated as non-LLM")
})
t.Run("malformed body", func(t *testing.T) {
_, err := p.ParseResponse(200, "application/json", []byte(`{not json`))
require.ErrorIs(t, err, ErrMalformedResponse, "bad JSON yields malformed error")
})
// Responses-API fixture: /v1/responses returns input_tokens/output_tokens
// (Anthropic-style) instead of prompt_tokens/completion_tokens. The parser
// must accept both.
t.Run("responses api fixture", func(t *testing.T) {
body, err := os.ReadFile(filepath.Join("fixtures", "openai_responses.json"))
require.NoError(t, err, "fixture must be readable")
usage, err := p.ParseResponse(200, "application/json", body)
require.NoError(t, err)
assert.Equal(t, int64(15), usage.InputTokens, "input_tokens should map directly")
assert.Equal(t, int64(414), usage.OutputTokens, "output_tokens should map directly")
assert.Equal(t, int64(429), usage.TotalTokens, "total_tokens carried through")
})
t.Run("responses api naming preferred over chat-completions when both present", func(t *testing.T) {
body := []byte(`{"usage":{"prompt_tokens":1,"completion_tokens":2,"input_tokens":15,"output_tokens":414,"total_tokens":429}}`)
usage, err := p.ParseResponse(200, "application/json", body)
require.NoError(t, err)
assert.Equal(t, int64(15), usage.InputTokens, "responses-api names take precedence")
assert.Equal(t, int64(414), usage.OutputTokens, "responses-api names take precedence")
})
t.Run("chat-completions naming still works alone", func(t *testing.T) {
body := []byte(`{"usage":{"prompt_tokens":15,"completion_tokens":414,"total_tokens":429}}`)
usage, err := p.ParseResponse(200, "application/json", body)
require.NoError(t, err)
assert.Equal(t, int64(15), usage.InputTokens, "prompt_tokens fallback")
assert.Equal(t, int64(414), usage.OutputTokens, "completion_tokens fallback")
})
// Cached-prompt accounting. cached_tokens is a SUBSET of
// prompt_tokens — input_tokens carries the full prompt count and
// the cached subset is reported separately so the cost meter can
// apply the discount rate to that portion.
t.Run("chat-completions cached_tokens subset surfaces", func(t *testing.T) {
body := []byte(`{"usage":{"prompt_tokens":1024,"completion_tokens":200,"total_tokens":1224,"prompt_tokens_details":{"cached_tokens":768}}}`)
usage, err := p.ParseResponse(200, "application/json", body)
require.NoError(t, err)
assert.Equal(t, int64(1024), usage.InputTokens, "input remains the full prompt count — cached is a subset, not a separate bucket")
assert.Equal(t, int64(768), usage.CachedInputTokens, "cached_tokens must surface so cost meter can discount the cached subset")
assert.Zero(t, usage.CacheCreationTokens, "OpenAI has no cache_creation analogue")
})
t.Run("responses-api input_tokens_details.cached_tokens surfaces", func(t *testing.T) {
body := []byte(`{"usage":{"input_tokens":2048,"output_tokens":100,"total_tokens":2148,"input_tokens_details":{"cached_tokens":1500}}}`)
usage, err := p.ParseResponse(200, "application/json", body)
require.NoError(t, err)
assert.Equal(t, int64(2048), usage.InputTokens)
assert.Equal(t, int64(1500), usage.CachedInputTokens, "Responses-API input_tokens_details.cached_tokens path must surface too")
})
t.Run("responses-api cached takes precedence over chat-completions when both present", func(t *testing.T) {
body := []byte(`{"usage":{"prompt_tokens":1,"input_tokens":2,"output_tokens":3,"prompt_tokens_details":{"cached_tokens":50},"input_tokens_details":{"cached_tokens":99}}}`)
usage, err := p.ParseResponse(200, "application/json", body)
require.NoError(t, err)
assert.Equal(t, int64(99), usage.CachedInputTokens, "Responses-API field wins when both naming conventions are present")
})
t.Run("absent cached_tokens leaves cached counts at zero", func(t *testing.T) {
body := []byte(`{"usage":{"prompt_tokens":15,"completion_tokens":414,"total_tokens":429}}`)
usage, err := p.ParseResponse(200, "application/json", body)
require.NoError(t, err)
assert.Zero(t, usage.CachedInputTokens, "no prompt_tokens_details = no cached subset")
})
}
func TestOpenAIExtractPrompt_ChatCompletions(t *testing.T) {
body := []byte(`{"model":"gpt-4o-mini","messages":[{"role":"system","content":"be brief"},{"role":"user","content":"ping"}]}`)
got := OpenAIParser{}.ExtractPrompt(body)
require.NotEmpty(t, got, "messages array must extract")
require.Contains(t, got, "system: be brief", "system role and content surface")
require.Contains(t, got, "user: ping", "user role and content surface")
}
func TestOpenAIExtractPrompt_ResponsesAPIStringInput(t *testing.T) {
body := []byte(`{"model":"gpt-5.4","input":"Hello there"}`)
got := OpenAIParser{}.ExtractPrompt(body)
require.Equal(t, "Hello there", got, "string input field should pass through")
}
func TestOpenAIExtractPrompt_ResponsesAPIInputParts(t *testing.T) {
body := []byte(`{"model":"gpt-5.4","input":[{"type":"input_text","input_text":"first"},{"type":"input_text","input_text":"second"}]}`)
got := OpenAIParser{}.ExtractPrompt(body)
require.Contains(t, got, "first", "first content part surfaces")
require.Contains(t, got, "second", "second content part surfaces")
}
// TestOpenAIExtractPrompt_ResponsesAPIMessageItems guards the live Codex
// shape: input is an array of message items whose text is nested under
// content[].text, not flat content parts. The old code fed the outer array
// to the content-part decoder and extracted nothing, so the stored prompt
// was empty.
func TestOpenAIExtractPrompt_ResponsesAPIMessageItems(t *testing.T) {
body := []byte(`{"model":"gpt-5.5","input":[` +
`{"type":"message","role":"developer","content":[{"type":"input_text","text":"system rules"}]},` +
`{"type":"message","role":"user","content":[{"type":"input_text","text":"hello there"}]},` +
`{"type":"reasoning","encrypted_content":"opaque","summary":[]},` +
`{"type":"message","role":"assistant","content":[{"type":"output_text","text":"prior reply"}]}` +
`]}`)
got := OpenAIParser{}.ExtractPrompt(body)
require.Contains(t, got, "system rules", "developer message content must surface")
require.Contains(t, got, "hello there", "user message content must surface")
require.Contains(t, got, "developer:", "role labels must prefix each message")
require.NotContains(t, got, "opaque", "reasoning items without text must be skipped")
}
func TestOpenAIExtractPrompt_LegacyCompletion(t *testing.T) {
body := []byte(`{"model":"text-davinci-003","prompt":"once upon a time"}`)
got := OpenAIParser{}.ExtractPrompt(body)
require.Equal(t, "once upon a time", got, "string prompt field should pass through")
}
func TestOpenAIExtractSessionID(t *testing.T) {
t.Run("codex client_metadata.session_id", func(t *testing.T) {
body := []byte(`{"model":"gpt-5.5","client_metadata":{"session_id":"019eeb72-ab7c-7cd2","thread_id":"t1"},"input":[]}`)
assert.Equal(t, "019eeb72-ab7c-7cd2", OpenAIParser{}.ExtractSessionID(body), "Codex session id must come from client_metadata.session_id")
})
t.Run("plain chat has no session", func(t *testing.T) {
body := []byte(`{"model":"gpt-4o","messages":[{"role":"user","content":"hi"}]}`)
assert.Equal(t, "", OpenAIParser{}.ExtractSessionID(body), "plain chat.completions carries no session id")
})
t.Run("non-JSON yields empty", func(t *testing.T) {
assert.Equal(t, "", OpenAIParser{}.ExtractSessionID([]byte("not json")), "malformed body must not error")
})
}
func TestOpenAIExtractCompletion_ChatCompletions(t *testing.T) {
body, err := os.ReadFile(filepath.Join("fixtures", "openai_chat_completion.json"))
require.NoError(t, err)
got := OpenAIParser{}.ExtractCompletion(200, "application/json", body)
require.NotEmpty(t, got, "fixture has assistant content")
}
func TestOpenAIExtractCompletion_ResponsesAPI(t *testing.T) {
body, err := os.ReadFile(filepath.Join("fixtures", "openai_responses.json"))
require.NoError(t, err)
got := OpenAIParser{}.ExtractCompletion(200, "application/json", body)
require.NotEmpty(t, got, "responses-api fixture has output content")
}
func TestOpenAIExtractCompletion_Streaming(t *testing.T) {
got := OpenAIParser{}.ExtractCompletion(200, "text/event-stream", []byte(""))
require.Empty(t, got, "streaming responses are skipped")
}
func TestOpenAIExtractCompletion_NonOK(t *testing.T) {
got := OpenAIParser{}.ExtractCompletion(500, "application/json", []byte(`{"choices":[{"message":{"content":"x"}}]}`))
require.Empty(t, got, "non-200 returns empty")
}
proxy/internal/llm/parser.go
+112
View File
@@ -0,0 +1,112 @@
// Package llm provides the shared LLM request and response parsing
// library consumed by proxy middleware. It is runtime agnostic: the same
// package is used by the native built-in executor now and will be reused
// by the WASM adapter later.
package llm
// Provider identifies an LLM API provider.
type Provider int
const (
// ProviderUnknown signals that no parser matched the request.
ProviderUnknown Provider = 0
// ProviderOpenAI identifies the OpenAI API surface.
ProviderOpenAI Provider = 1
// ProviderAnthropic identifies the Anthropic Messages API surface.
ProviderAnthropic Provider = 2
// ProviderBedrock identifies the AWS Bedrock runtime surface.
ProviderBedrock Provider = 3
)
// RequestFacts captures the subset of the LLM request body that the
// middleware annotates as metadata (model, streaming flag). Additional
// fields are added as parsers grow.
type RequestFacts struct {
Model string
Stream bool
}
// Usage is the provider-agnostic token accounting emitted to metrics and
// access logs. Downstream consumers map InputTokens/OutputTokens to the
// plg.llm.* metadata allowlist entries.
//
// CachedInputTokens carries OpenAI's prompt_tokens_details.cached_tokens
// (a SUBSET of InputTokens) when the response is from OpenAI, or
// Anthropic's cache_read_input_tokens (ADDITIVE to InputTokens) when from
// Anthropic. The cost meter switches formula on KeyLLMProvider so the
// two shapes are billed correctly without double-counting.
//
// CacheCreationTokens carries Anthropic's cache_creation_input_tokens
// (ADDITIVE; not present in the OpenAI shape).
type Usage struct {
InputTokens int64
OutputTokens int64
TotalTokens int64
CachedInputTokens int64
CacheCreationTokens int64
}
// Parser is the per-provider interface implemented in this package. The
// dispatcher selects a parser by calling DetectFromURL against the incoming
// request path; ties break by registration order (see Parsers).
type Parser interface {
Provider() Provider
ProviderName() string
DetectFromURL(path string) bool
ParseRequest(body []byte) (RequestFacts, error)
ParseResponse(status int, contentType string, body []byte) (Usage, error)
// ExtractPrompt returns the user-facing prompt text from a request body.
// Different endpoint shapes (chat.completions, responses, messages) are
// handled by the per-provider implementation. Returns "" when no prompt
// can be extracted; never returns an error — extraction is best-effort
// because callers use the result for observability, not authorization.
ExtractPrompt(body []byte) string
// ExtractCompletion returns the assistant-facing completion text from a
// non-streaming response body. status and contentType match the
// ParseResponse arguments so implementations can fast-fail uniformly.
ExtractCompletion(status int, contentType string, body []byte) string
// ExtractSessionID returns a stable identifier that groups requests of
// the same conversation / coding session, read from the per-provider
// location clients populate (e.g. OpenAI Codex's client_metadata.session_id,
// Claude Code's metadata.user_id). Returns "" when the body carries no
// recognised session marker; extraction is best-effort and never errors.
ExtractSessionID(body []byte) string
}
// Parsers returns the built-in parser set in a stable order. The order is
// deterministic so that DetectFromURL ties produce consistent routing.
func Parsers() []Parser {
return []Parser{
OpenAIParser{},
AnthropicParser{},
BedrockParser{},
}
}
// DetectParser returns the first parser whose DetectFromURL matches the given
// request path. ok=false means no parser claimed the path.
func DetectParser(path string) (Parser, bool) {
for _, p := range Parsers() {
if p.DetectFromURL(path) {
return p, true
}
}
return nil, false
}
// ParserByName returns the parser whose ProviderName matches id. Used by
// callers that already know which provider surface a request will hit
// (e.g. the agent-network middleware chain configured per synthesised
// service) so they can skip URL sniffing. ok=false when no parser is
// registered under that name.
func ParserByName(id string) (Parser, bool) {
if id == "" {
return nil, false
}
for _, p := range Parsers() {
if p.ProviderName() == id {
return p, true
}
}
return nil, false
}
proxy/internal/llm/parser_test.go
+54
View File
@@ -0,0 +1,54 @@
package llm
import (
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestParsers_ProviderNames(t *testing.T) {
parsers := Parsers()
require.Len(t, parsers, 3, "three built-in parsers expected")
names := make([]string, 0, len(parsers))
for _, p := range parsers {
names = append(names, p.ProviderName())
}
assert.Contains(t, names, "openai", "OpenAI parser should be registered")
assert.Contains(t, names, "anthropic", "Anthropic parser should be registered")
assert.Contains(t, names, "bedrock", "Bedrock parser should be registered")
}
func TestDetectParser(t *testing.T) {
cases := []struct {
name string
path string
expectedName string
expectOK bool
}{
{"openai chat", "/v1/chat/completions", "openai", true},
{"openai prefixed", "/api/v1/chat/completions", "openai", true},
{"openai responses", "/v1/responses", "openai", true},
{"anthropic messages", "/v1/messages", "anthropic", true},
{"anthropic prefixed", "/proxy/v1/messages?query", "anthropic", true},
{"unknown path", "/healthz", "", false},
{"empty path", "", "", false},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
p, ok := DetectParser(tc.path)
require.Equal(t, tc.expectOK, ok, "detection success mismatch for %q", tc.path)
if ok {
assert.Equal(t, tc.expectedName, p.ProviderName(), "provider name mismatch")
}
})
}
}
func TestProviderValues(t *testing.T) {
assert.Equal(t, Provider(0), ProviderUnknown, "unknown provider is the zero value")
assert.Equal(t, ProviderOpenAI, OpenAIParser{}.Provider(), "OpenAI parser reports its provider enum")
assert.Equal(t, ProviderAnthropic, AnthropicParser{}.Provider(), "Anthropic parser reports its provider enum")
}
proxy/internal/llm/pricing/defaults_coverage_test.go
+65
View File
@@ -0,0 +1,65 @@
package pricing
import (
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
// TestDefaultTable_FirstPartyModelCoverage guards the embedded defaults against
// silent drift/gaps: every metered first-party model the management catalog
// enumerates must resolve to a price, and a few rates that previously drifted
// are pinned to their LiteLLM-validated values. Keep this list in step with the
// catalog (management/server/agentnetwork/catalog) when adding models.
func TestDefaultTable_FirstPartyModelCoverage(t *testing.T) {
tbl := DefaultTable()
require.NotNil(t, tbl, "embedded default pricing table must load")
mustPrice := map[string][]string{
// openai parser covers openai_api, azure_openai_api, and mistral_api.
"openai": {
"gpt-5.5", "gpt-5.5-pro", "gpt-5.4", "gpt-5.4-mini", "gpt-5.4-nano",
"gpt-5.3-codex", "gpt-5.3-chat-latest", "o4-mini",
"gpt-4.1", "gpt-4.1-mini", "gpt-4.1-nano", "gpt-4o", "gpt-4o-mini",
"gpt-4-turbo", "gpt-3.5-turbo", "gpt-35-turbo",
"text-embedding-3-large", "text-embedding-3-small",
"mistral-large-latest", "mistral-medium-3-5", "codestral-2508",
"ministral-8b-latest", "mistral-embed",
},
"anthropic": {
"claude-fable-5", "claude-opus-5", "claude-opus-4-8", "claude-opus-4-7", "claude-opus-4-6",
"claude-opus-4-1", "claude-sonnet-4-6", "claude-sonnet-4-5", "claude-haiku-4-5",
},
// bedrock keys are the normalized ids the request parser emits.
"bedrock": {
"anthropic.claude-opus-5", "anthropic.claude-opus-4-8", "anthropic.claude-opus-4-7", "anthropic.claude-opus-4-6",
"anthropic.claude-opus-4-1", "anthropic.claude-sonnet-4-6", "anthropic.claude-sonnet-4-5",
"anthropic.claude-haiku-4-5", "meta.llama3-3-70b-instruct",
"amazon.nova-pro", "amazon.nova-lite", "amazon.nova-micro", "amazon.nova-2-lite",
},
}
for provider, models := range mustPrice {
for _, m := range models {
_, ok := tbl.Cost(provider, m, 1000, 1000, 0, 0)
assert.True(t, ok, "%s/%s must be priced in the embedded defaults", provider, m)
}
}
// Pin per-direction rates independently (input-only then output-only) so a
// swap or skew of input<->output that preserves the combined total is still
// caught — these are rates that previously drifted or are easy to mis-enter.
in, ok := tbl.Cost("openai", "gpt-5.4", 1000, 0, 0, 0)
require.True(t, ok)
assert.InDelta(t, 0.0025, in, 1e-9, "gpt-5.4 input = 0.0025 per 1k")
out, ok := tbl.Cost("openai", "gpt-5.4", 0, 1000, 0, 0)
require.True(t, ok)
assert.InDelta(t, 0.015, out, 1e-9, "gpt-5.4 output = 0.015 per 1k")
in, ok = tbl.Cost("bedrock", "anthropic.claude-sonnet-4-5", 1000, 0, 0, 0)
require.True(t, ok)
assert.InDelta(t, 0.003, in, 1e-9, "bedrock sonnet-4-5 input = 0.003 per 1k")
out, ok = tbl.Cost("bedrock", "anthropic.claude-sonnet-4-5", 0, 1000, 0, 0)
require.True(t, ok)
assert.InDelta(t, 0.015, out, 1e-9, "bedrock sonnet-4-5 output = 0.015 per 1k")
}
proxy/internal/llm/pricing/defaults_pricing.yaml
+298
View File
@@ -0,0 +1,298 @@
# Embedded default pricing for llm_observability. Compiled into the proxy
# binary via go:embed in pricing.go; cost annotation works out of the box
# without any operator action.
#
# Operators override entries by dropping a pricing.yaml into --plugin-data-dir
# (or whichever basename is given via params.pricing_path). The override file
# only needs entries the operator wants to change; missing entries fall
# through to these defaults.
#
# Values are USD per 1_000 tokens. Public list prices drift; ship a fresh
# binary or override individual entries via the override file as needed.
#
# Optional cache fields:
# cached_input_per_1k OpenAI: rate for prompt_tokens_details.cached_tokens
# (a SUBSET of prompt_tokens). Typically 0.5x input.
# Absent → cached portion bills at input_per_1k.
# cache_read_per_1k Anthropic: rate for cache_read_input_tokens
# (ADDITIVE to input_tokens). Typically 0.1x input.
# Absent → cache reads bill at input_per_1k.
# cache_creation_per_1k Anthropic: rate for cache_creation_input_tokens
# (ADDITIVE to input_tokens). Typically 1.25x input.
# Absent → cache writes bill at input_per_1k.
openai:
# OpenAI + OpenAI-compatible providers (openai_api, azure_openai_api,
# mistral_api, and the openai-parser gateways) all emit llm.provider="openai",
# so their models are priced here. Kept in sync with the management catalog;
# rates cross-checked against LiteLLM model_prices_and_context_window.json.
# GPT-5.x family — cache reads 10% of input (0.1x).
gpt-5.5:
input_per_1k: 0.005
output_per_1k: 0.03
cached_input_per_1k: 0.0005
gpt-5.5-pro:
input_per_1k: 0.03
output_per_1k: 0.18
cached_input_per_1k: 0.003
gpt-5.4:
input_per_1k: 0.0025
output_per_1k: 0.015
cached_input_per_1k: 0.00025
gpt-5.4-pro:
input_per_1k: 0.03
output_per_1k: 0.18
cached_input_per_1k: 0.003
gpt-5.4-mini:
input_per_1k: 0.00075
output_per_1k: 0.0045
cached_input_per_1k: 0.000075
gpt-5.4-nano:
input_per_1k: 0.0002
output_per_1k: 0.00125
cached_input_per_1k: 0.00002
gpt-5.3-codex:
input_per_1k: 0.00175
output_per_1k: 0.014
cached_input_per_1k: 0.000175
gpt-5.3-chat-latest:
input_per_1k: 0.00175
output_per_1k: 0.014
cached_input_per_1k: 0.000175
# GPT-5 (2025) family — kept for gateway requests using the unsuffixed ids.
gpt-5:
input_per_1k: 0.00125
output_per_1k: 0.01
cached_input_per_1k: 0.000125
gpt-5-mini:
input_per_1k: 0.00025
output_per_1k: 0.002
cached_input_per_1k: 0.000025
gpt-5-nano:
input_per_1k: 0.00005
output_per_1k: 0.0004
cached_input_per_1k: 0.000005
o4-mini:
input_per_1k: 0.0011
output_per_1k: 0.0044
cached_input_per_1k: 0.000275
# GPT-4.1 family — cache reads 25% of input.
gpt-4.1:
input_per_1k: 0.002
output_per_1k: 0.008
cached_input_per_1k: 0.0005
gpt-4.1-mini:
input_per_1k: 0.0004
output_per_1k: 0.0016
cached_input_per_1k: 0.0001
gpt-4.1-nano:
input_per_1k: 0.0001
output_per_1k: 0.0004
cached_input_per_1k: 0.000025
# GPT-4o family — cache reads 50% of input (0.5x).
gpt-4o:
input_per_1k: 0.0025
output_per_1k: 0.01
cached_input_per_1k: 0.00125
gpt-4o-mini:
input_per_1k: 0.00015
output_per_1k: 0.0006
cached_input_per_1k: 0.000075
# Older GPT — no prompt caching.
gpt-4-turbo:
input_per_1k: 0.01
output_per_1k: 0.03
gpt-3.5-turbo:
input_per_1k: 0.0005
output_per_1k: 0.0015
gpt-35-turbo: # Azure deployment alias of gpt-3.5-turbo
input_per_1k: 0.0005
output_per_1k: 0.0015
# Embeddings — no caching, no output tokens.
text-embedding-3-large:
input_per_1k: 0.00013
output_per_1k: 0
text-embedding-3-small:
input_per_1k: 0.00002
output_per_1k: 0
# Mistral (mistral_api) — routed via the openai parser; no prompt caching.
mistral-large-latest:
input_per_1k: 0.0005
output_per_1k: 0.0015
mistral-medium-latest:
input_per_1k: 0.0004
output_per_1k: 0.002
mistral-medium-3-5:
input_per_1k: 0.0015
output_per_1k: 0.0075
mistral-small-latest:
input_per_1k: 0.00006
output_per_1k: 0.00018
magistral-medium-latest:
input_per_1k: 0.002
output_per_1k: 0.005
magistral-small-latest:
input_per_1k: 0.0005
output_per_1k: 0.0015
devstral-medium-latest:
input_per_1k: 0.0004
output_per_1k: 0.002
devstral-small-latest:
input_per_1k: 0.0001
output_per_1k: 0.0003
codestral-2508:
input_per_1k: 0.0003
output_per_1k: 0.0009
codestral-latest:
input_per_1k: 0.001
output_per_1k: 0.003
ministral-3-14b-2512:
input_per_1k: 0.0002
output_per_1k: 0.0002
ministral-8b-latest:
input_per_1k: 0.00015
output_per_1k: 0.00015
ministral-3-3b-2512:
input_per_1k: 0.0001
output_per_1k: 0.0001
mistral-embed:
input_per_1k: 0.0001
output_per_1k: 0
# Kimi / Moonshot AI (kimi_api) — OpenAI-compatible /v1 endpoint. Moonshot
# reports cache hits OpenAI-style when present; cached input is 10% of
# input ($0.30 vs $3.00 per MTok). kimi-k3 is the only model the platform
# serves newer accounts (K2-era ids and kimi-latest 404), matching the
# management catalog.
kimi-k3:
input_per_1k: 0.003
output_per_1k: 0.015
cached_input_per_1k: 0.0003
anthropic:
# Claude 4.x family — cache reads ≈10% of input, cache writes ≈125% of input.
# Pricing source: Anthropic's current published rates per million tokens,
# divided by 1000 for the per-1k figures stored here.
claude-fable-5:
input_per_1k: 0.010
output_per_1k: 0.050
cache_read_per_1k: 0.001
cache_creation_per_1k: 0.0125
claude-opus-5:
input_per_1k: 0.005
output_per_1k: 0.025
cache_read_per_1k: 0.0005
cache_creation_per_1k: 0.00625
claude-opus-4-8:
input_per_1k: 0.005
output_per_1k: 0.025
cache_read_per_1k: 0.0005
cache_creation_per_1k: 0.00625
claude-opus-4-7:
input_per_1k: 0.005
output_per_1k: 0.025
cache_read_per_1k: 0.0005
cache_creation_per_1k: 0.00625
claude-opus-4-6:
input_per_1k: 0.005
output_per_1k: 0.025
cache_read_per_1k: 0.0005
cache_creation_per_1k: 0.00625
claude-opus-4-1:
input_per_1k: 0.015
output_per_1k: 0.075
cache_read_per_1k: 0.0015
cache_creation_per_1k: 0.01875
claude-sonnet-4-6:
input_per_1k: 0.003
output_per_1k: 0.015
cache_read_per_1k: 0.0003
cache_creation_per_1k: 0.00375
claude-sonnet-4-5:
input_per_1k: 0.003
output_per_1k: 0.015
cache_read_per_1k: 0.0003
cache_creation_per_1k: 0.00375
claude-haiku-4-5:
input_per_1k: 0.001
output_per_1k: 0.005
cache_read_per_1k: 0.0001
cache_creation_per_1k: 0.00125
# Kimi / Moonshot AI (kimi_api) via the Anthropic-compatible endpoint
# (/anthropic/v1/messages — the official Claude Code setup). Same rates
# as the OpenAI-shape entry above. "kimi-k3[1m]" is the model id some
# Claude Code guides set for the 1M-context alias; priced identically so
# cost metering doesn't silently skip those requests.
kimi-k3:
input_per_1k: 0.003
output_per_1k: 0.015
cache_read_per_1k: 0.0003
"kimi-k3[1m]":
input_per_1k: 0.003
output_per_1k: 0.015
cache_read_per_1k: 0.0003
bedrock:
# AWS Bedrock model ids, normalised by the request parser (cross-region
# inference-profile prefix + version/throughput suffix stripped), e.g.
# eu.anthropic.claude-sonnet-4-5-20250929-v1:0 -> anthropic.claude-sonnet-4-5.
# Anthropic-on-Bedrock keeps the additive cache buckets (read ≈0.1x input,
# write ≈1.25x input); Nova / Llama report no cache, so cost is input+output.
anthropic.claude-opus-5:
input_per_1k: 0.005
output_per_1k: 0.025
cache_read_per_1k: 0.0005
cache_creation_per_1k: 0.00625
anthropic.claude-opus-4-8:
input_per_1k: 0.005
output_per_1k: 0.025
cache_read_per_1k: 0.0005
cache_creation_per_1k: 0.00625
anthropic.claude-opus-4-7:
input_per_1k: 0.005
output_per_1k: 0.025
cache_read_per_1k: 0.0005
cache_creation_per_1k: 0.00625
anthropic.claude-opus-4-6:
input_per_1k: 0.005
output_per_1k: 0.025
cache_read_per_1k: 0.0005
cache_creation_per_1k: 0.00625
anthropic.claude-opus-4-1:
input_per_1k: 0.015
output_per_1k: 0.075
cache_read_per_1k: 0.0015
cache_creation_per_1k: 0.01875
anthropic.claude-sonnet-4-6:
input_per_1k: 0.003
output_per_1k: 0.015
cache_read_per_1k: 0.0003
cache_creation_per_1k: 0.00375
anthropic.claude-sonnet-4-5:
input_per_1k: 0.003
output_per_1k: 0.015
cache_read_per_1k: 0.0003
cache_creation_per_1k: 0.00375
anthropic.claude-haiku-4-5:
input_per_1k: 0.001
output_per_1k: 0.005
cache_read_per_1k: 0.0001
cache_creation_per_1k: 0.00125
meta.llama3-3-70b-instruct:
input_per_1k: 0.00072
output_per_1k: 0.00072
amazon.nova-2-lite:
input_per_1k: 0.0003
output_per_1k: 0.0025
amazon.nova-pro:
input_per_1k: 0.0008
output_per_1k: 0.0032
amazon.nova-lite:
input_per_1k: 0.00006
output_per_1k: 0.00024
amazon.nova-micro:
input_per_1k: 0.000035
output_per_1k: 0.00014
proxy/internal/llm/pricing/pricing.go
+489
View File
@@ -0,0 +1,489 @@
// Package pricing implements the embedded-default + override pricing table
// shared by middleware that converts LLM token usage into a USD cost
// estimate. The table is hot-reloadable from a basename under the proxy
// data directory; missing override files keep the embedded defaults so
// cost annotation works without operator action.
package pricing
import (
"bytes"
"context"
_ "embed"
"errors"
"fmt"
"io"
"io/fs"
"math"
"path/filepath"
"regexp"
"strings"
"sync"
"sync/atomic"
"time"
log "github.com/sirupsen/logrus"
"go.opentelemetry.io/otel/attribute"
"go.opentelemetry.io/otel/metric"
"gopkg.in/yaml.v3"
)
//go:embed defaults_pricing.yaml
var defaultPricingYAML []byte
var (
defaultTableOnce sync.Once
defaultTablePtr *Table
)
// DefaultTable returns the pricing table embedded in the binary. The result
// is parsed once and shared; callers must not mutate the returned value.
// Cost annotation works without any operator action because every loader
// starts with this table.
func DefaultTable() *Table {
defaultTableOnce.Do(func() {
t, err := parsePricingBytes(defaultPricingYAML)
if err != nil {
panic(fmt.Sprintf("llmobs: embedded default pricing failed to parse: %v", err))
}
defaultTablePtr = t
})
return defaultTablePtr
}
// mergeOver returns a new Table containing every entry from base, with any
// matching entry from overlay replacing the base value. Either argument may
// be nil. Result is a fresh allocation so callers can mutate / Store safely.
func mergeOver(base, overlay *Table) *Table {
if overlay == nil || len(overlay.entries) == 0 {
return base
}
if base == nil || len(base.entries) == 0 {
return overlay
}
out := make(map[string]map[string]Entry, len(base.entries))
for provider, models := range base.entries {
inner := make(map[string]Entry, len(models))
for model, e := range models {
inner[model] = e
}
out[provider] = inner
}
for provider, models := range overlay.entries {
inner, ok := out[provider]
if !ok {
inner = make(map[string]Entry, len(models))
out[provider] = inner
}
for model, e := range models {
inner[model] = e
}
}
return &Table{entries: out}
}
// Entry is a single model's input and output pricing, expressed in USD per
// 1000 tokens.
//
// CachedInputPer1K applies to OpenAI's cached prompt tokens, which are a
// subset of input_tokens — when set, the cached portion is billed at this
// rate and the non-cached remainder at InputPer1K. Zero means "no discount
// configured", and cached tokens are billed at InputPer1K (matches current
// behaviour where cached counts weren't extracted at all).
//
// CacheReadPer1K and CacheCreationPer1K apply to Anthropic's two prompt-
// cache fields, which are additive to input_tokens: cache_read is the
// cheaper read-from-cache rate, cache_creation is the more expensive
// write-to-cache rate. Zero means "no rate configured" and the
// corresponding token bucket is billed at InputPer1K. This is more
// accurate than today's behaviour, where Anthropic's cache tokens are
// ignored and not charged at all.
type Entry struct {
InputPer1K float64
OutputPer1K float64
CachedInputPer1K float64
CacheReadPer1K float64
CacheCreationPer1K float64
}
// Table is a provider-to-model pricing lookup. Instances are immutable once
// built and are swapped atomically by Loader.
type Table struct {
entries map[string]map[string]Entry
}
// Cost returns the estimated USD cost for the given token counts. ok is
// false when the provider or model is not present in the table; the caller
// can still emit token metrics with a model=unknown label.
//
// Provider-shape semantics for cached / cache-creation counts:
//
// - OpenAI: cachedInput is a SUBSET of inTokens. The cached portion is
// billed at CachedInputPer1K (or InputPer1K when no override), and the
// non-cached remainder of inTokens at InputPer1K. cacheCreation is
// ignored (OpenAI has no analogue).
// - Anthropic: cachedInput (cache_read) and cacheCreation are ADDITIVE to
// inTokens. The three buckets are billed at CacheReadPer1K,
// CacheCreationPer1K, and InputPer1K respectively, each falling back
// to InputPer1K when the corresponding rate is zero.
// - Other providers: cached and cacheCreation are ignored; cost is
// inTokens*InputPer1K + outTokens*OutputPer1K.
func (t *Table) Cost(provider, model string, inTokens, outTokens, cachedInput, cacheCreation int64) (float64, bool) {
c, ok := t.Costs(provider, model, inTokens, outTokens, cachedInput, cacheCreation)
return c.TotalUSD, ok
}
// Costs is a per-request cost split. The four per-bucket fields are the base
// of the breakdown — one per token bucket the provider bills separately — and
// the two aggregates are derived from them:
//
// TotalUSD = InputUSD + CachedInputUSD + CacheCreationUSD + OutputUSD
// CacheUSD = CachedInputUSD + CacheCreationUSD
//
// InputUSD is always the cost of the *non-cached* input bucket, for both
// provider shapes: on OpenAI the cached subset is carved out of inTokens and
// billed as CachedInputUSD, so the two never double-count. Buckets a provider
// doesn't bill are zero, which keeps the identities above true everywhere.
type Costs struct {
InputUSD float64
CachedInputUSD float64
CacheCreationUSD float64
OutputUSD float64
TotalUSD float64
CacheUSD float64
}
// newCosts assembles a split from its per-bucket parts, deriving the two
// aggregates so TotalUSD and CacheUSD can never drift from the breakdown.
func newCosts(input, cachedInput, cacheCreation, output float64) Costs {
return Costs{
InputUSD: input,
CachedInputUSD: cachedInput,
CacheCreationUSD: cacheCreation,
OutputUSD: output,
TotalUSD: input + cachedInput + cacheCreation + output,
CacheUSD: cachedInput + cacheCreation,
}
}
// Costs returns the estimated USD cost split for the given token counts, with
// the same semantics as Cost.
func (t *Table) Costs(provider, model string, inTokens, outTokens, cachedInput, cacheCreation int64) (Costs, bool) {
// Clamp negatives to zero before any pricing math so a malformed
// upstream count can never produce a negative cost.
if inTokens < 0 {
inTokens = 0
}
if outTokens < 0 {
outTokens = 0
}
if cachedInput < 0 {
cachedInput = 0
}
if cacheCreation < 0 {
cacheCreation = 0
}
if t == nil {
return Costs{}, false
}
byModel, ok := t.entries[provider]
if !ok {
return Costs{}, false
}
entry, ok := byModel[model]
if !ok {
return Costs{}, false
}
output := (float64(outTokens) / 1000.0) * entry.OutputPer1K
switch provider {
case "openai":
// cachedInput is a subset of inTokens; clamp so a malformed
// upstream (cached > total) can't produce a negative remainder.
clamped := cachedInput
if clamped > inTokens {
clamped = inTokens
}
cachedRate := entry.CachedInputPer1K
if cachedRate <= 0 {
cachedRate = entry.InputPer1K
}
nonCached := float64(inTokens-clamped) / 1000.0 * entry.InputPer1K
cached := float64(clamped) / 1000.0 * cachedRate
return newCosts(nonCached, cached, 0, output), true
case "anthropic", "bedrock":
// Bedrock-Anthropic returns the same additive cache buckets as
// first-party Anthropic; non-Anthropic Bedrock models simply report
// zero cache tokens, so this formula degrades to input + output.
readRate := entry.CacheReadPer1K
if readRate <= 0 {
readRate = entry.InputPer1K
}
createRate := entry.CacheCreationPer1K
if createRate <= 0 {
createRate = entry.InputPer1K
}
input := float64(inTokens) / 1000.0 * entry.InputPer1K
read := float64(cachedInput) / 1000.0 * readRate
create := float64(cacheCreation) / 1000.0 * createRate
return newCosts(input, read, create, output), true
default:
input := float64(inTokens) / 1000.0 * entry.InputPer1K
return newCosts(input, 0, 0, output), true
}
}
// Has reports whether the provider/model pair is present in the table.
func (t *Table) Has(provider, model string) bool {
if t == nil {
return false
}
byModel, ok := t.entries[provider]
if !ok {
return false
}
_, ok = byModel[model]
return ok
}
// pricingFile mirrors the on-disk YAML schema. Keys are provider names; the
// nested map keys are model names.
type pricingFile map[string]map[string]struct {
InputPer1K float64 `yaml:"input_per_1k"`
OutputPer1K float64 `yaml:"output_per_1k"`
CachedInputPer1K float64 `yaml:"cached_input_per_1k"`
CacheReadPer1K float64 `yaml:"cache_read_per_1k"`
CacheCreationPer1K float64 `yaml:"cache_creation_per_1k"`
}
const (
// ReloadInterval is the mtime-poll cadence for the background reloader.
ReloadInterval = 30 * time.Second
// errorBackoff bounds how often the loader logs a repeated parse error.
errorBackoff = 5 * time.Minute
)
var basenameRegex = regexp.MustCompile(`^[a-zA-Z0-9._-]+$`)
// Loader is a confined, hot-reloadable pricing table reader. Construction
// must succeed against the target file; subsequent reload failures keep the
// previously-loaded table so callers never observe a blank price list.
type Loader struct {
baseDir string
fullPath string
pluginID string
table atomic.Pointer[Table]
mtime atomic.Int64
failures metric.Int64Counter
interval time.Duration
}
// NewLoader returns a pricing loader that overlays an optional file-based
// table on top of the embedded defaults. Missing override file, baseDir, or
// relPath is not an error: the loader keeps the embedded defaults so cost
// metadata is still emitted for known models.
//
// Errors:
// - bad basename, traversal segment, or absolute relPath are rejected so a
// misconfigured target surfaces immediately.
// - permission errors and YAML parse errors keep the defaults but log a
// warning; cost annotation does not silently break.
//
// failures is optional; pass nil in tests that do not care about
// reload-failure telemetry.
func NewLoader(baseDir, relPath, pluginID string, failures metric.Int64Counter) (*Loader, error) {
defaults := DefaultTable()
l := &Loader{
baseDir: baseDir,
pluginID: pluginID,
failures: failures,
}
l.table.Store(defaults)
if strings.TrimSpace(baseDir) == "" || strings.TrimSpace(relPath) == "" {
return l, nil
}
full, err := resolveMiddlewareDataPath(baseDir, relPath)
if err != nil {
return nil, err
}
l.fullPath = full
overlay, mtime, err := loadPricing(full)
if err != nil {
if errors.Is(err, fs.ErrNotExist) {
// Override file is optional. Defaults already stored.
return l, nil
}
// Symlink rejection, oversize file, parse failure, permission errors
// — surface so a misconfigured operator sees the problem instead of
// silently running with stale defaults.
return nil, fmt.Errorf("load pricing %s: %w", full, err)
}
l.table.Store(mergeOver(defaults, overlay))
l.mtime.Store(mtime.UnixNano())
return l, nil
}
// Get returns the current pricing table. The returned pointer is immutable;
// callers must not mutate its contents.
func (l *Loader) Get() *Table {
if l == nil {
return nil
}
return l.table.Load()
}
// WatchesFile reports whether this loader is bound to an override file on
// disk. False for defaults-only loaders (no operator override given).
// Callers use this to decide whether to spawn the mtime-poll goroutine.
func (l *Loader) WatchesFile() bool {
if l == nil {
return false
}
return l.fullPath != ""
}
// SetReloadInterval overrides the mtime-poll cadence used by Reload. Calls
// after Reload has started have no effect on the running loop. Intended for
// tests; production code uses the default ReloadInterval.
func (l *Loader) SetReloadInterval(d time.Duration) {
if l == nil || d <= 0 {
return
}
l.interval = d
}
// Reload runs a polling loop that checks the pricing file mtime every
// ReloadInterval (or the value passed to SetReloadInterval). Returns when
// ctx is cancelled.
func (l *Loader) Reload(ctx context.Context) {
if l == nil {
return
}
interval := l.interval
if interval <= 0 {
interval = ReloadInterval
}
t := time.NewTicker(interval)
defer t.Stop()
var lastErrAt time.Time
for {
select {
case <-ctx.Done():
return
case <-t.C:
if err := l.reload(); err != nil {
if l.failures != nil {
l.failures.Add(ctx, 1, metric.WithAttributes(
attribute.String("plugin", l.pluginID),
))
}
now := time.Now()
if now.Sub(lastErrAt) >= errorBackoff {
log.Warnf("llmobs: pricing reload failed for %s: %v", l.fullPath, err)
lastErrAt = now
}
}
}
}
}
// reload performs a single-shot mtime check and reload. The reloaded
// override file is merged on top of the embedded defaults; missing override
// (e.g. operator deleted the file) is not an error and reverts to defaults.
func (l *Loader) reload() error {
if l.fullPath == "" {
// Defaults-only loader; nothing on disk to reload.
return nil
}
mtime, err := statMtime(l.fullPath)
if err != nil {
if errors.Is(err, fs.ErrNotExist) {
// File was removed since startup. Drop back to defaults and
// reset mtime so a future re-creation triggers a reload.
l.table.Store(DefaultTable())
l.mtime.Store(0)
return nil
}
return err
}
if mtime.UnixNano() == l.mtime.Load() {
return nil
}
overlay, newMtime, err := loadPricing(l.fullPath)
if err != nil {
return err
}
l.table.Store(mergeOver(DefaultTable(), overlay))
l.mtime.Store(newMtime.UnixNano())
return nil
}
// resolveMiddlewareDataPath validates relPath is a safe basename and resolves
// it under baseDir. An additional cleaned-prefix check guards against
// CVE-style edge cases where Join is used with trailing path segments.
func resolveMiddlewareDataPath(baseDir, relPath string) (string, error) {
if strings.TrimSpace(baseDir) == "" {
return "", errors.New("middleware-data-dir is not configured")
}
if relPath == "" {
return "", errors.New("pricing path is empty")
}
if !basenameRegex.MatchString(relPath) {
return "", fmt.Errorf("pricing path %q is not a safe basename", relPath)
}
if filepath.IsAbs(relPath) {
return "", fmt.Errorf("pricing path %q must be a basename, not absolute", relPath)
}
cleanBase, err := filepath.Abs(filepath.Clean(baseDir))
if err != nil {
return "", fmt.Errorf("resolve middleware-data-dir: %w", err)
}
full := filepath.Join(cleanBase, relPath)
cleanedFull := filepath.Clean(full)
if !strings.HasPrefix(cleanedFull, cleanBase+string(filepath.Separator)) && cleanedFull != cleanBase {
return "", fmt.Errorf("pricing path %q escapes middleware-data-dir", relPath)
}
return cleanedFull, nil
}
func parsePricingBytes(data []byte) (*Table, error) {
dec := yaml.NewDecoder(bytes.NewReader(data))
dec.KnownFields(true)
var raw pricingFile
if err := dec.Decode(&raw); err != nil && !errors.Is(err, io.EOF) {
return nil, fmt.Errorf("decode pricing yaml: %w", err)
}
out := make(map[string]map[string]Entry, len(raw))
for provider, models := range raw {
inner := make(map[string]Entry, len(models))
for model, entry := range models {
for field, v := range map[string]float64{
"input_per_1k": entry.InputPer1K,
"output_per_1k": entry.OutputPer1K,
"cached_input_per_1k": entry.CachedInputPer1K,
"cache_read_per_1k": entry.CacheReadPer1K,
"cache_creation_per_1k": entry.CacheCreationPer1K,
} {
if v < 0 || math.IsNaN(v) || math.IsInf(v, 0) {
return nil, fmt.Errorf("pricing %s/%s: %s must be a finite, non-negative rate, got %v", provider, model, field, v)
}
}
inner[model] = Entry{
InputPer1K: entry.InputPer1K,
OutputPer1K: entry.OutputPer1K,
CachedInputPer1K: entry.CachedInputPer1K,
CacheReadPer1K: entry.CacheReadPer1K,
CacheCreationPer1K: entry.CacheCreationPer1K,
}
}
out[provider] = inner
}
return &Table{entries: out}, nil
}
proxy/internal/llm/pricing/pricing_other.go
+20
View File
@@ -0,0 +1,20 @@
//go:build !unix
package pricing
import (
"fmt"
"time"
)
// loadPricing is unavailable on non-Unix platforms because O_NOFOLLOW and
// fstat-from-FD are required to honour the spec's symlink-safety rules. The
// proxy is only deployed on Linux today; a Windows port would need an
// equivalent path-as-handle implementation.
func loadPricing(path string) (*Table, time.Time, error) {
return nil, time.Time{}, fmt.Errorf("llmobs pricing loader is not supported on this platform: %s", path)
}
func statMtime(path string) (time.Time, error) {
return time.Time{}, fmt.Errorf("llmobs pricing loader is not supported on this platform: %s", path)
}
proxy/internal/llm/pricing/pricing_test.go
+432
View File
@@ -0,0 +1,432 @@
//go:build unix
package pricing
import (
"context"
"os"
"path/filepath"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func copyFixture(t *testing.T, src, dst string) {
t.Helper()
data, err := os.ReadFile(src)
require.NoError(t, err, "read source fixture")
require.NoError(t, os.WriteFile(dst, data, 0o600), "write target fixture")
}
func TestNewLoader_HappyPath(t *testing.T) {
base := t.TempDir()
copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), filepath.Join(base, "pricing.yaml"))
l, err := NewLoader(base, "pricing.yaml", "llm_observability", nil)
require.NoError(t, err, "NewLoader must succeed with a valid fixture")
table := l.Get()
require.NotNil(t, table, "table populated after load")
cost, ok := table.Cost("openai", "gpt-4o-mini", 1000, 1000, 0, 0)
require.True(t, ok, "known provider/model resolves")
assert.InDelta(t, 0.00075, cost, 1e-9, "cost = 0.00015 + 0.0006 per 1k tokens")
cost, ok = table.Cost("openai", "gpt-4o", 2000, 1000, 0, 0)
require.True(t, ok, "second known model resolves")
assert.InDelta(t, 0.015, cost, 1e-9, "cost for gpt-4o: 2*0.0025 + 1*0.01")
cost, ok = table.Cost("anthropic", "claude-sonnet-4-5", 1000, 1000, 0, 0)
require.True(t, ok, "anthropic model resolves")
assert.InDelta(t, 0.018, cost, 1e-9, "cost for claude-sonnet-4-5: 0.003 + 0.015")
}
// TestCost_OpenAICachedSubsetDiscount proves OpenAI's cached input
// tokens are billed at the configured cached_input_per_1k rate while
// the non-cached remainder of input_tokens is billed at the regular
// rate. Critical because OpenAI returns cached_tokens as a SUBSET of
// prompt_tokens — naïvely charging the cached count on top of
// prompt_tokens would double-bill that portion.
func TestCost_OpenAICachedSubsetDiscount(t *testing.T) {
tbl := &Table{entries: map[string]map[string]Entry{
"openai": {"gpt-4o": {
InputPer1K: 0.0025, // 0.0025 USD per 1k input tokens
OutputPer1K: 0.01,
CachedInputPer1K: 0.00125, // 0.5x discount on cached
}},
}}
// 1000 prompt tokens, 750 of which were cached. 250 non-cached
// at regular rate, 750 cached at the discount rate, 500 output.
cost, ok := tbl.Cost("openai", "gpt-4o", 1000, 500, 750, 0)
require.True(t, ok, "known model resolves")
want := (250.0/1000.0)*0.0025 + (750.0/1000.0)*0.00125 + (500.0/1000.0)*0.01
assert.InDelta(t, want, cost, 1e-12,
"cached subset must bill at the discount rate; non-cached remainder at regular rate")
}
// TestCost_OpenAICachedFallsBackToInputRate covers the operator
// opt-in contract: when CachedInputPer1K is unset (zero), cached
// tokens bill at the regular input rate. This matches today's
// behaviour (cached counts weren't extracted at all so they
// implicitly billed at the input rate via prompt_tokens).
func TestCost_OpenAICachedFallsBackToInputRate(t *testing.T) {
tbl := &Table{entries: map[string]map[string]Entry{
"openai": {"gpt-4o": {InputPer1K: 0.0025, OutputPer1K: 0.01}},
}}
cost, ok := tbl.Cost("openai", "gpt-4o", 1000, 500, 750, 0)
require.True(t, ok)
want := 0.0025 + (500.0/1000.0)*0.01
assert.InDelta(t, want, cost, 1e-12,
"absent cached_input_per_1k rate must fall back to input_per_1k — same as pre-feature behaviour")
}
// TestCost_OpenAIClampsCachedToInputCount is the defensive guard
// against malformed upstream responses that report cached_tokens >
// prompt_tokens. We clamp so the formula never produces a negative
// "non-cached remainder" multiplied by the input rate.
func TestCost_OpenAIClampsCachedToInputCount(t *testing.T) {
tbl := &Table{entries: map[string]map[string]Entry{
"openai": {"gpt-4o": {InputPer1K: 0.0025, OutputPer1K: 0.01, CachedInputPer1K: 0.00125}},
}}
cost, ok := tbl.Cost("openai", "gpt-4o", 100, 0, 9999, 0)
require.True(t, ok)
// All 100 cached, 0 non-cached. Output is 0.
want := (100.0 / 1000.0) * 0.00125
assert.InDelta(t, want, cost, 1e-12,
"cached count > input count must clamp to input — never bill negative non-cached tokens")
}
// TestCost_AnthropicCacheReadAndCreationAreAdditive proves the
// Anthropic shape: cache_read and cache_creation tokens are
// ADDITIVE to input_tokens (not subset), each billed at its own
// configured rate. The two rates pull in opposite directions —
// cache_read is the cheaper read-from-cache rate (≈0.1× input),
// cache_creation is the more expensive write-to-cache rate
// (≈1.25× input).
func TestCost_AnthropicCacheReadAndCreationAreAdditive(t *testing.T) {
tbl := &Table{entries: map[string]map[string]Entry{
"anthropic": {"claude-sonnet": {
InputPer1K: 0.003,
OutputPer1K: 0.015,
CacheReadPer1K: 0.0003, // 0.1x of input
CacheCreationPer1K: 0.00375, // 1.25x of input
}},
}}
// 256 regular input + 768 cache_read + 512 cache_creation +
// 200 output. Each input bucket bills at its own rate.
cost, ok := tbl.Cost("anthropic", "claude-sonnet", 256, 200, 768, 512)
require.True(t, ok, "known model resolves")
want := (256.0/1000.0)*0.003 +
(768.0/1000.0)*0.0003 +
(512.0/1000.0)*0.00375 +
(200.0/1000.0)*0.015
assert.InDelta(t, want, cost, 1e-12,
"each Anthropic input bucket must bill at its own configured rate")
}
// TestCost_AnthropicCacheRatesFallBackToInput covers the no-opt-in
// path: when neither CacheReadPer1K nor CacheCreationPer1K is set,
// cache tokens bill at the regular input rate. This is more
// accurate than today's behaviour (cache tokens ignored entirely)
// without requiring operators to opt in via YAML.
func TestCost_AnthropicCacheRatesFallBackToInput(t *testing.T) {
tbl := &Table{entries: map[string]map[string]Entry{
"anthropic": {"claude-sonnet": {InputPer1K: 0.003, OutputPer1K: 0.015}},
}}
cost, ok := tbl.Cost("anthropic", "claude-sonnet", 256, 200, 768, 512)
require.True(t, ok)
// Without overrides: every input bucket at input_per_1k.
want := ((256.0+768.0+512.0)/1000.0)*0.003 + (200.0/1000.0)*0.015
assert.InDelta(t, want, cost, 1e-12,
"absent cache rates must fall back to input_per_1k — Anthropic cache tokens were ignored before this change, billing at input rate is more accurate as a default")
}
func TestNewLoader_UnknownModel(t *testing.T) {
base := t.TempDir()
copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), filepath.Join(base, "pricing.yaml"))
l, err := NewLoader(base, "pricing.yaml", "llm_observability", nil)
require.NoError(t, err)
_, ok := l.Get().Cost("openai", "fantasy-model", 10, 10, 0, 0)
assert.False(t, ok, "unknown model returns ok=false")
_, ok = l.Get().Cost("cohere", "anything", 10, 10, 0, 0)
assert.False(t, ok, "unknown provider returns ok=false")
}
func TestNewLoader_InvalidYAMLRejected(t *testing.T) {
base := t.TempDir()
require.NoError(t, os.WriteFile(filepath.Join(base, "pricing.yaml"), []byte("\t- this is not: valid: yaml: :["), 0o600))
_, err := NewLoader(base, "pricing.yaml", "llm_observability", nil)
require.Error(t, err, "invalid YAML must surface as construction error")
}
func TestLoader_ReloadKeepsPreviousOnParseError(t *testing.T) {
base := t.TempDir()
target := filepath.Join(base, "pricing.yaml")
copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), target)
l, err := NewLoader(base, "pricing.yaml", "llm_observability", nil)
require.NoError(t, err)
require.NotNil(t, l.Get(), "initial table populated")
// Overwrite with content that violates the strict schema (extra field)
// plus a bumped mtime to trigger reload.
require.NoError(t, os.WriteFile(target, []byte("openai:\n gpt-4o:\n input_per_1k: 1.0\n output_per_1k: 2.0\n bogus_field: nope\n"), 0o600))
future := time.Now().Add(time.Hour)
require.NoError(t, os.Chtimes(target, future, future))
err = l.reload()
require.Error(t, err, "parse error surfaced by reload()")
cost, ok := l.Get().Cost("openai", "gpt-4o-mini", 1000, 1000, 0, 0)
require.True(t, ok, "previous table still available after parse failure")
assert.InDelta(t, 0.00075, cost, 1e-9, "previous cost preserved")
}
func TestLoader_ReloadNoChangeIsNoOp(t *testing.T) {
base := t.TempDir()
target := filepath.Join(base, "pricing.yaml")
copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), target)
l, err := NewLoader(base, "pricing.yaml", "llm_observability", nil)
require.NoError(t, err)
ptrBefore := l.Get()
require.NoError(t, l.reload(), "no-change reload must not error")
ptrAfter := l.Get()
assert.Same(t, ptrBefore, ptrAfter, "table pointer unchanged when mtime unchanged")
}
func TestLoader_ReloadDetectsChange(t *testing.T) {
base := t.TempDir()
target := filepath.Join(base, "pricing.yaml")
copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), target)
l, err := NewLoader(base, "pricing.yaml", "llm_observability", nil)
require.NoError(t, err)
updated := []byte("openai:\n gpt-4o-mini:\n input_per_1k: 1.00\n output_per_1k: 2.00\n")
require.NoError(t, os.WriteFile(target, updated, 0o600))
future := time.Now().Add(time.Hour)
require.NoError(t, os.Chtimes(target, future, future))
require.NoError(t, l.reload(), "reload must succeed on valid new content")
cost, ok := l.Get().Cost("openai", "gpt-4o-mini", 1000, 1000, 0, 0)
require.True(t, ok, "updated model still present")
assert.InDelta(t, 3.0, cost, 0.0001, "new prices are applied: 1 + 2 per 1k")
}
// TestLoader_ReloadGoroutinePicksUpChanges proves the background goroutine
// started via Reload actually swaps the pricing table when the file changes
// on disk. Without that goroutine running, pricing edits would never reach
// requests until a proxy restart.
func TestLoader_ReloadGoroutinePicksUpChanges(t *testing.T) {
base := t.TempDir()
target := filepath.Join(base, "pricing.yaml")
copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), target)
l, err := NewLoader(base, "pricing.yaml", "llm_observability", nil)
require.NoError(t, err)
l.SetReloadInterval(20 * time.Millisecond)
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
done := make(chan struct{})
go func() {
l.Reload(ctx)
close(done)
}()
// Before any rewrite, the loader holds the fixture's prices.
costBefore, ok := l.Get().Cost("openai", "gpt-4o-mini", 1000, 1000, 0, 0)
require.True(t, ok, "fixture model must resolve initially")
assert.InDelta(t, 0.00075, costBefore, 1e-9, "fixture prices apply before rewrite")
updated := []byte("openai:\n gpt-4o-mini:\n input_per_1k: 1.00\n output_per_1k: 2.00\n")
require.NoError(t, os.WriteFile(target, updated, 0o600))
future := time.Now().Add(time.Hour)
require.NoError(t, os.Chtimes(target, future, future))
deadline := time.Now().Add(2 * time.Second)
for {
cost, ok := l.Get().Cost("openai", "gpt-4o-mini", 1000, 1000, 0, 0)
if ok && cost > 2.5 {
break
}
if time.Now().After(deadline) {
t.Fatalf("background reloader did not pick up rewrite within deadline")
}
time.Sleep(10 * time.Millisecond)
}
cancel()
select {
case <-done:
case <-time.After(2 * time.Second):
t.Fatal("Reload loop did not exit after cancel")
}
}
func TestLoader_ReloadBackgroundLoopCancellation(t *testing.T) {
base := t.TempDir()
copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), filepath.Join(base, "pricing.yaml"))
l, err := NewLoader(base, "pricing.yaml", "llm_observability", nil)
require.NoError(t, err)
ctx, cancel := context.WithCancel(context.Background())
done := make(chan struct{})
go func() {
l.Reload(ctx)
close(done)
}()
cancel()
select {
case <-done:
case <-time.After(2 * time.Second):
t.Fatal("Reload loop did not exit on context cancel")
}
}
func TestNewLoader_PathValidation(t *testing.T) {
base := t.TempDir()
copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), filepath.Join(base, "pricing.yaml"))
cases := []struct {
name string
relPath string
}{
{"traversal", "../../etc/passwd"},
{"absolute", "/etc/passwd"},
{"slash in basename", "sub/pricing.yaml"},
{"control chars", "pricing\x00.yaml"},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
_, err := NewLoader(base, tc.relPath, "llm_observability", nil)
require.Error(t, err, "NewLoader must reject %q", tc.relPath)
})
}
// Empty relPath is no longer a validation error: the loader treats it
// as "no override file, defaults only" so cost metadata is still
// emitted for the embedded models out of the box.
t.Run("empty falls back to defaults", func(t *testing.T) {
l, err := NewLoader(base, "", "llm_observability", nil)
require.NoError(t, err, "empty relPath should yield a defaults-only loader")
require.NotNil(t, l, "loader must be returned")
require.False(t, l.WatchesFile(), "no file watching when no override is given")
_, ok := l.Get().Cost("openai", "gpt-4o-mini", 1000, 1000, 0, 0)
assert.True(t, ok, "embedded defaults should still resolve gpt-4o-mini")
})
}
// TestNewLoader_PathValidation_Extended covers the remaining attack shapes
// called out in C2: dot references, embedded traversal segments, and a
// newline in the basename. The basename regex must reject each one even
// though filepath.Clean would otherwise collapse them.
func TestNewLoader_PathValidation_Extended(t *testing.T) {
base := t.TempDir()
copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), filepath.Join(base, "pricing.yaml"))
cases := []struct {
name string
relPath string
}{
{"dot", "."},
{"dotdot", ".."},
{"relative traversal", "../pricing.yaml"},
{"embedded slash", "pri/cing.yaml"},
{"newline", "pricing\n.yaml"},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
_, err := NewLoader(base, tc.relPath, "llm_observability", nil)
require.Error(t, err, "NewLoader must reject %q", tc.relPath)
})
}
}
// TestNewLoader_ValidBasenameLoads proves the allowlist is exclusive: a
// basename containing only safe characters under baseDir loads. Without this
// a regression that over-tightened the regex would silently break valid
// deployments.
func TestNewLoader_ValidBasenameLoads(t *testing.T) {
base := t.TempDir()
copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), filepath.Join(base, "pricing-v2_prod.yaml"))
l, err := NewLoader(base, "pricing-v2_prod.yaml", "llm_observability", nil)
require.NoError(t, err, "basename with _, -, . must load")
require.NotNil(t, l.Get(), "table populated")
}
// TestNewLoader_SymlinkOutsideBaseDirRejected constructs a symlink under
// baseDir that points to a file outside it. O_NOFOLLOW must refuse to open
// the symlink even though the symlink path itself is a valid basename under
// baseDir.
func TestNewLoader_SymlinkOutsideBaseDirRejected(t *testing.T) {
outside := t.TempDir()
target := filepath.Join(outside, "evil.yaml")
copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), target)
base := t.TempDir()
link := filepath.Join(base, "pricing.yaml")
require.NoError(t, os.Symlink(target, link), "symlink setup")
_, err := NewLoader(base, "pricing.yaml", "llm_observability", nil)
require.Error(t, err, "O_NOFOLLOW must reject symlink even when it points outside baseDir")
}
func TestNewLoader_SymlinkRejected(t *testing.T) {
base := t.TempDir()
concrete := filepath.Join(base, "real.yaml")
copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), concrete)
link := filepath.Join(base, "pricing.yaml")
require.NoError(t, os.Symlink(concrete, link), "symlink setup")
_, err := NewLoader(base, "pricing.yaml", "llm_observability", nil)
require.Error(t, err, "O_NOFOLLOW must reject symlinked targets")
}
func TestTableCost_NilSafe(t *testing.T) {
var t1 *Table
cost, ok := t1.Cost("x", "y", 1, 1, 0, 0)
assert.False(t, ok, "nil table reports unknown")
assert.Zero(t, cost, "nil table returns zero cost")
assert.False(t, t1.Has("x", "y"), "nil table has nothing")
}
func TestLoaderGet_NilSafe(t *testing.T) {
var l *Loader
assert.Nil(t, l.Get(), "nil loader returns nil table")
}
// TestNewLoader_RejectsOversizedFile_FixesM4 proves the loader bounds reads
// at maxPricingBytes so a hostile file cannot exhaust process memory.
func TestNewLoader_RejectsOversizedFile_FixesM4(t *testing.T) {
base := t.TempDir()
target := filepath.Join(base, "pricing.yaml")
// Build a YAML payload larger than the cap. We pad with valid YAML
// comments so a partial read would still fail the size check rather
// than the parser.
header := "openai:\n"
bigComment := make([]byte, maxPricingBytes+1024)
for i := range bigComment {
bigComment[i] = ' '
}
bigComment[0] = '#'
bigComment[len(bigComment)-1] = '\n'
payload := append([]byte(header), bigComment...)
require.NoError(t, os.WriteFile(target, payload, 0o600))
_, err := NewLoader(base, "pricing.yaml", "llm_observability", nil)
require.Error(t, err, "oversized pricing file must be rejected")
assert.Contains(t, err.Error(), "exceeds", "rejection must reference the byte cap")
}
proxy/internal/llm/pricing/pricing_unix.go
+68
View File
@@ -0,0 +1,68 @@
//go:build unix
package pricing
import (
"fmt"
"io"
"os"
"syscall"
"time"
log "github.com/sirupsen/logrus"
)
// maxPricingBytes caps the size of the pricing YAML on read so a hostile or
// runaway file cannot exhaust process memory during reload. 1 MiB is several
// orders of magnitude larger than any reasonable pricing table.
const maxPricingBytes int64 = 1 << 20
// loadPricing opens the file with O_NOFOLLOW, fstats the open descriptor,
// and parses from that same descriptor. Never re-opens by path so a
// mid-read rename or symlink swap cannot substitute content. Bytes are
// capped at maxPricingBytes so the loader cannot be coerced into reading an
// unbounded file.
func loadPricing(path string) (*Table, time.Time, error) {
f, err := os.OpenFile(path, os.O_RDONLY|syscall.O_NOFOLLOW, 0)
if err != nil {
return nil, time.Time{}, fmt.Errorf("open %s: %w", path, err)
}
defer func() {
if cerr := f.Close(); cerr != nil {
log.Debugf("close pricing file %s: %v", path, cerr)
}
}()
info, err := f.Stat()
if err != nil {
return nil, time.Time{}, fmt.Errorf("fstat %s: %w", path, err)
}
if !info.Mode().IsRegular() {
return nil, time.Time{}, fmt.Errorf("pricing file %s is not a regular file", path)
}
data, err := io.ReadAll(io.LimitReader(f, maxPricingBytes+1))
if err != nil {
return nil, time.Time{}, fmt.Errorf("read %s: %w", path, err)
}
if int64(len(data)) > maxPricingBytes {
return nil, time.Time{}, fmt.Errorf("pricing file %s exceeds %d bytes", path, maxPricingBytes)
}
table, err := parsePricingBytes(data)
if err != nil {
return nil, time.Time{}, err
}
return table, info.ModTime(), nil
}
// statMtime returns the mtime of the file at path. It uses lstat semantics
// via os.Lstat so a symlink swap is detected even though O_NOFOLLOW will
// later reject the open.
func statMtime(path string) (time.Time, error) {
info, err := os.Lstat(path)
if err != nil {
return time.Time{}, fmt.Errorf("lstat %s: %w", path, err)
}
return info.ModTime(), nil
}
proxy/internal/llm/sse.go
+117
View File
@@ -0,0 +1,117 @@
package llm
import (
"bufio"
"errors"
"fmt"
"io"
"strings"
)
// Event represents a single server-sent event. Type is the dispatch name
// carried on an "event:" line (empty when the stream uses only "data:"
// lines). Data is the concatenation of every "data:" line that made up the
// event, joined by a single newline.
type Event struct {
Type string
Data string
}
// Scanner reads SSE events from an underlying byte stream. Events are
// delimited by a blank line ("\n\n"). CRLF line endings are normalized to LF
// transparently so fixtures captured from live servers can be replayed.
//
// Scanner is not safe for concurrent use.
type Scanner struct {
r *bufio.Reader
maxLine int
}
// NewScanner wraps the given reader. The default underlying buffer size is
// large enough for typical provider events (~64 KiB); callers needing
// larger events can wrap the reader in their own bufio.Reader beforehand.
func NewScanner(r io.Reader) *Scanner {
return &Scanner{
r: bufio.NewReaderSize(r, 64*1024),
maxLine: 1 << 20,
}
}
// Next returns the next event. It returns io.EOF after the final event has
// been consumed. A trailing event that is not terminated by a blank line is
// still returned before io.EOF so that servers which close the connection
// without a trailing newline are handled correctly.
func (s *Scanner) Next() (Event, error) {
var (
event Event
dataBuf strings.Builder
hasData bool
hasAny bool
)
for {
line, err := s.readLine()
if err != nil {
if errors.Is(err, io.EOF) && hasAny {
event.Data = dataBuf.String()
return event, nil
}
return Event{}, err
}
if line == "" {
if !hasAny {
continue
}
event.Data = dataBuf.String()
return event, nil
}
hasAny = true
if strings.HasPrefix(line, ":") {
continue
}
field, value := splitField(line)
switch field {
case "event":
event.Type = value
case "data":
if hasData {
dataBuf.WriteByte('\n')
}
dataBuf.WriteString(value)
hasData = true
}
}
}
func (s *Scanner) readLine() (string, error) {
line, err := s.r.ReadString('\n')
if err != nil {
if errors.Is(err, io.EOF) && line != "" {
return trimEOL(line), nil
}
return "", err
}
if len(line) > s.maxLine {
return "", fmt.Errorf("sse line exceeds %d bytes", s.maxLine)
}
return trimEOL(line), nil
}
func trimEOL(line string) string {
line = strings.TrimRight(line, "\n")
line = strings.TrimRight(line, "\r")
return line
}
func splitField(line string) (string, string) {
idx := strings.IndexByte(line, ':')
if idx < 0 {
return line, ""
}
field := line[:idx]
value := strings.TrimPrefix(line[idx+1:], " ")
return field, value
}
proxy/internal/llm/sse_test.go
+175
View File
@@ -0,0 +1,175 @@
package llm
import (
"errors"
"io"
"os"
"path/filepath"
"strings"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func collectEvents(t *testing.T, r io.Reader) []Event {
t.Helper()
s := NewScanner(r)
var out []Event
for {
ev, err := s.Next()
if errors.Is(err, io.EOF) {
return out
}
require.NoError(t, err, "unexpected error scanning SSE")
out = append(out, ev)
}
}
func TestSSEScanner_OpenAIFixture(t *testing.T) {
f, err := os.Open(filepath.Join("fixtures", "openai_stream.txt"))
require.NoError(t, err, "fixture must be openable")
defer f.Close()
events := collectEvents(t, f)
require.Len(t, events, 4, "expected 4 data frames (3 chunks + [DONE])")
for _, ev := range events {
assert.Empty(t, ev.Type, "OpenAI stream uses data-only frames")
}
assert.Contains(t, events[2].Data, `"usage"`, "third chunk carries usage block")
assert.Equal(t, "[DONE]", events[3].Data, "final frame is the OpenAI DONE sentinel")
}
func TestSSEScanner_AnthropicFixture(t *testing.T) {
f, err := os.Open(filepath.Join("fixtures", "anthropic_stream.txt"))
require.NoError(t, err, "fixture must be openable")
defer f.Close()
events := collectEvents(t, f)
require.Len(t, events, 7, "expected 7 Anthropic events")
types := make([]string, 0, len(events))
for _, ev := range events {
types = append(types, ev.Type)
}
assert.Equal(t, []string{
"message_start",
"content_block_start",
"content_block_delta",
"content_block_delta",
"content_block_stop",
"message_delta",
"message_stop",
}, types, "Anthropic event ordering matches fixture")
var deltaUsage Event
for _, ev := range events {
if ev.Type == "message_delta" {
deltaUsage = ev
break
}
}
assert.Contains(t, deltaUsage.Data, `"output_tokens":45`, "message_delta carries partial usage")
}
func TestSSEScanner_MultilineData(t *testing.T) {
raw := "event: ping\ndata: line1\ndata: line2\ndata: line3\n\n"
events := collectEvents(t, strings.NewReader(raw))
require.Len(t, events, 1, "one logical event from three data lines")
assert.Equal(t, "ping", events[0].Type, "event name honored")
assert.Equal(t, "line1\nline2\nline3", events[0].Data, "data lines joined with newline")
}
func TestSSEScanner_CRLF(t *testing.T) {
raw := "event: foo\r\ndata: bar\r\n\r\ndata: baz\r\n\r\n"
events := collectEvents(t, strings.NewReader(raw))
require.Len(t, events, 2, "CRLF-delimited events recognized")
assert.Equal(t, "foo", events[0].Type, "first event type preserved")
assert.Equal(t, "bar", events[0].Data, "first event data preserved")
assert.Empty(t, events[1].Type, "second event has no event name")
assert.Equal(t, "baz", events[1].Data, "second event data preserved")
}
func TestSSEScanner_EmptyInput(t *testing.T) {
s := NewScanner(strings.NewReader(""))
_, err := s.Next()
require.ErrorIs(t, err, io.EOF, "empty input yields immediate EOF")
}
func TestSSEScanner_CommentIgnored(t *testing.T) {
raw := ": this is a comment\ndata: hi\n\n"
events := collectEvents(t, strings.NewReader(raw))
require.Len(t, events, 1, "comment line does not emit an event")
assert.Equal(t, "hi", events[0].Data, "data line honoured after comment")
}
func TestSSEScanner_TrailingWithoutBlankLine(t *testing.T) {
raw := "event: foo\ndata: bar\n"
events := collectEvents(t, strings.NewReader(raw))
require.Len(t, events, 1, "trailing event without blank line still emitted")
assert.Equal(t, "foo", events[0].Type)
assert.Equal(t, "bar", events[0].Data)
}
// TestSSEScanner_ManyConsecutiveEmptyLines feeds a stream that is nothing
// but empty lines. The scanner must terminate without panic — empty lines
// alone do not constitute an event and must yield io.EOF.
func TestSSEScanner_ManyConsecutiveEmptyLines(t *testing.T) {
raw := strings.Repeat("\n", 100)
s := NewScanner(strings.NewReader(raw))
_, err := s.Next()
require.ErrorIs(t, err, io.EOF, "100 empty lines must terminate as EOF without panic")
}
// TestSSEScanner_InterleavedCRLFAndLF mixes \r\n and \n terminators within
// the same event. The scanner normalizes both and must still recover a
// coherent event.
func TestSSEScanner_InterleavedCRLFAndLF(t *testing.T) {
raw := "event: mix\r\ndata: first\ndata: second\r\n\n"
events := collectEvents(t, strings.NewReader(raw))
require.Len(t, events, 1, "mixed line endings must still produce one event")
assert.Equal(t, "mix", events[0].Type)
assert.Equal(t, "first\nsecond", events[0].Data, "both data lines joined")
}
// TestSSEScanner_LongSingleDataLine constructs a single data line that
// exceeds the default bufio buffer (64 KiB) but stays under the scanner
// maxLine. The scanner must round-trip the value intact without panicking
// or truncating silently.
func TestSSEScanner_LongSingleDataLine(t *testing.T) {
big := strings.Repeat("x", 80<<10)
raw := "data: " + big + "\n\n"
events := collectEvents(t, strings.NewReader(raw))
require.Len(t, events, 1, "long single-line event must be emitted")
assert.Equal(t, big, events[0].Data, "long data preserved")
}
// TestSSEScanner_BinaryGarbageInData validates that non-printable bytes
// inside a data line do not crash the parser. The scanner should either
// round-trip them or return a well-formed error — never panic.
func TestSSEScanner_BinaryGarbageInData(t *testing.T) {
raw := "data: \x00\x01\x02\xff\xfe\n\n"
defer func() {
if r := recover(); r != nil {
t.Fatalf("scanner panicked on binary garbage: %v", r)
}
}()
s := NewScanner(strings.NewReader(raw))
ev, err := s.Next()
require.NoError(t, err, "binary bytes in data should not surface as error")
assert.Equal(t, "\x00\x01\x02\xff\xfe", ev.Data, "binary payload round-trips")
}
// TestSSEScanner_UnknownFieldsIgnored stresses the field parser by sending
// unrecognized field names ("id:", "retry:", "custom:"). They must be
// silently ignored per the SSE spec; the scanner must not panic or emit
// spurious events.
func TestSSEScanner_UnknownFieldsIgnored(t *testing.T) {
raw := "id: 1\nretry: 5000\ncustom: value\ndata: payload\n\n"
events := collectEvents(t, strings.NewReader(raw))
require.Len(t, events, 1, "unknown fields must not spawn extra events")
assert.Equal(t, "payload", events[0].Data, "data field survives amid unknown fields")
}
proxy/internal/metrics/l4_metrics_test.go
+69
View File
@@ -0,0 +1,69 @@
package metrics_test
import (
"context"
"reflect"
"testing"
"time"
promexporter "go.opentelemetry.io/otel/exporters/prometheus"
sdkmetric "go.opentelemetry.io/otel/sdk/metric"
"github.com/netbirdio/netbird/proxy/internal/metrics"
"github.com/netbirdio/netbird/proxy/internal/types"
)
func newTestMetrics(t *testing.T) *metrics.Metrics {
t.Helper()
exporter, err := promexporter.New()
if err != nil {
t.Fatalf("create prometheus exporter: %v", err)
}
provider := sdkmetric.NewMeterProvider(sdkmetric.WithReader(exporter))
pkg := reflect.TypeOf(metrics.Metrics{}).PkgPath()
meter := provider.Meter(pkg)
m, err := metrics.New(context.Background(), meter)
if err != nil {
t.Fatalf("create metrics: %v", err)
}
return m
}
func TestL4ServiceGauge(t *testing.T) {
m := newTestMetrics(t)
m.L4ServiceAdded(types.ServiceModeTCP)
m.L4ServiceAdded(types.ServiceModeTCP)
m.L4ServiceAdded(types.ServiceModeUDP)
m.L4ServiceRemoved(types.ServiceModeTCP)
}
func TestTCPRelayMetrics(t *testing.T) {
m := newTestMetrics(t)
acct := types.AccountID("acct-1")
m.TCPRelayStarted(acct)
m.TCPRelayStarted(acct)
m.TCPRelayEnded(acct, 10*time.Second, 1000, 500)
m.TCPRelayDialError(acct)
m.TCPRelayRejected(acct)
}
func TestUDPSessionMetrics(t *testing.T) {
m := newTestMetrics(t)
acct := types.AccountID("acct-2")
m.UDPSessionStarted(acct)
m.UDPSessionStarted(acct)
m.UDPSessionEnded(acct)
m.UDPSessionDialError(acct)
m.UDPSessionRejected(acct)
m.UDPPacketRelayed(types.RelayDirectionClientToBackend, 100)
m.UDPPacketRelayed(types.RelayDirectionClientToBackend, 200)
m.UDPPacketRelayed(types.RelayDirectionBackendToClient, 150)
}
proxy/internal/metrics/metrics.go
+459
View File
@@ -0,0 +1,459 @@
package metrics
import (
"context"
"net/http"
"sync"
"time"
"go.opentelemetry.io/otel/attribute"
"go.opentelemetry.io/otel/metric"
"github.com/netbirdio/netbird/proxy/internal/proxy"
"github.com/netbirdio/netbird/proxy/internal/responsewriter"
"github.com/netbirdio/netbird/proxy/internal/types"
)
// Metrics collects OpenTelemetry metrics for the proxy.
type Metrics struct {
ctx context.Context
meter metric.Meter
requestsTotal metric.Int64Counter
activeRequests metric.Int64UpDownCounter
configuredDomains metric.Int64UpDownCounter
totalPaths metric.Int64UpDownCounter
requestDuration metric.Int64Histogram
backendDuration metric.Int64Histogram
certificateIssueDuration metric.Int64Histogram
// Management sync metrics.
snapshotSyncDuration metric.Int64Histogram
snapshotBatchDuration metric.Int64Histogram
addPeerDuration metric.Int64Histogram
// L4 service-level metrics.
l4Services metric.Int64UpDownCounter
// L4 TCP connection-level metrics.
tcpActiveConns metric.Int64UpDownCounter
tcpConnsTotal metric.Int64Counter
tcpConnDuration metric.Int64Histogram
tcpBytesTotal metric.Int64Counter
// L4 UDP session-level metrics.
udpActiveSess metric.Int64UpDownCounter
udpSessionsTotal metric.Int64Counter
udpPacketsTotal metric.Int64Counter
udpBytesTotal metric.Int64Counter
mappingsMux sync.Mutex
mappingPaths map[string]int
}
// Meter returns the OpenTelemetry meter the bundle was built with, so other
// subsystems (e.g. the middleware manager) register instruments on the same
// meter.
func (m *Metrics) Meter() metric.Meter {
return m.meter
}
// New creates a Metrics instance using the given OpenTelemetry meter.
func New(ctx context.Context, meter metric.Meter) (*Metrics, error) {
m := &Metrics{
ctx: ctx,
meter: meter,
mappingPaths: make(map[string]int),
}
if err := m.initHTTPMetrics(meter); err != nil {
return nil, err
}
if err := m.initSyncMetrics(meter); err != nil {
return nil, err
}
if err := m.initL4Metrics(meter); err != nil {
return nil, err
}
return m, nil
}
func (m *Metrics) initHTTPMetrics(meter metric.Meter) error {
var err error
m.requestsTotal, err = meter.Int64Counter(
"proxy.http.request.counter",
metric.WithUnit("1"),
metric.WithDescription("Total number of requests made to the netbird proxy"),
)
if err != nil {
return err
}
m.activeRequests, err = meter.Int64UpDownCounter(
"proxy.http.active_requests",
metric.WithUnit("1"),
metric.WithDescription("Current in-flight requests handled by the netbird proxy"),
)
if err != nil {
return err
}
m.configuredDomains, err = meter.Int64UpDownCounter(
"proxy.domains.count",
metric.WithUnit("1"),
metric.WithDescription("Current number of domains configured on the netbird proxy"),
)
if err != nil {
return err
}
m.totalPaths, err = meter.Int64UpDownCounter(
"proxy.paths.count",
metric.WithUnit("1"),
metric.WithDescription("Total number of paths configured on the netbird proxy"),
)
if err != nil {
return err
}
m.requestDuration, err = meter.Int64Histogram(
"proxy.http.request.duration.ms",
metric.WithUnit("milliseconds"),
metric.WithDescription("Duration of requests made to the netbird proxy"),
)
if err != nil {
return err
}
m.backendDuration, err = meter.Int64Histogram(
"proxy.backend.duration.ms",
metric.WithUnit("milliseconds"),
metric.WithDescription("Duration of peer round trip time from the netbird proxy"),
)
if err != nil {
return err
}
m.certificateIssueDuration, err = meter.Int64Histogram(
"proxy.certificate.issue.duration.ms",
metric.WithUnit("milliseconds"),
metric.WithDescription("Duration of ACME certificate issuance"),
)
return err
}
func (m *Metrics) initSyncMetrics(meter metric.Meter) error {
var err error
m.snapshotSyncDuration, err = meter.Int64Histogram(
"proxy.sync.snapshot.duration.ms",
metric.WithUnit("milliseconds"),
metric.WithDescription("Duration from management connect until the initial snapshot sync is complete"),
metric.WithExplicitBucketBoundaries(100, 250, 500, 1000, 2500, 5000, 10000, 30000, 60000, 120000, 300000),
)
if err != nil {
return err
}
m.snapshotBatchDuration, err = meter.Int64Histogram(
"proxy.sync.batch.duration.ms",
metric.WithUnit("milliseconds"),
metric.WithDescription("Duration to process a single mapping batch during initial snapshot sync"),
metric.WithExplicitBucketBoundaries(100, 250, 500, 1000, 2500, 5000, 10000, 30000, 60000, 120000, 300000),
)
if err != nil {
return err
}
m.addPeerDuration, err = meter.Int64Histogram(
"proxy.peer.add.duration.ms",
metric.WithUnit("milliseconds"),
metric.WithDescription("Duration to add a peer for an account (keygen + gRPC CreateProxyPeer + embed.New)"),
metric.WithExplicitBucketBoundaries(10, 25, 50, 100, 250, 500, 1000, 2500, 5000, 10000),
)
return err
}
// RecordSnapshotSyncDuration records the total time from connect to sync-complete.
func (m *Metrics) RecordSnapshotSyncDuration(d time.Duration) {
m.snapshotSyncDuration.Record(m.ctx, d.Milliseconds())
}
// RecordSnapshotBatchDuration records the time to process one mapping batch during initial sync.
func (m *Metrics) RecordSnapshotBatchDuration(d time.Duration) {
m.snapshotBatchDuration.Record(m.ctx, d.Milliseconds())
}
// RecordAddPeerDuration records the time to create a new peer for an account.
func (m *Metrics) RecordAddPeerDuration(d time.Duration, err error) {
result := "success"
if err != nil {
result = "error"
}
m.addPeerDuration.Record(m.ctx, d.Milliseconds(), metric.WithAttributes(
attribute.String("result", result),
))
}
func (m *Metrics) initL4Metrics(meter metric.Meter) error {
var err error
m.l4Services, err = meter.Int64UpDownCounter(
"proxy.l4.services.count",
metric.WithUnit("1"),
metric.WithDescription("Current number of configured L4 services (TCP/TLS/UDP) by mode"),
)
if err != nil {
return err
}
m.tcpActiveConns, err = meter.Int64UpDownCounter(
"proxy.tcp.active_connections",
metric.WithUnit("1"),
metric.WithDescription("Current number of active TCP/TLS relay connections"),
)
if err != nil {
return err
}
m.tcpConnsTotal, err = meter.Int64Counter(
"proxy.tcp.connections.total",
metric.WithUnit("1"),
metric.WithDescription("Total TCP/TLS relay connections by result and account"),
)
if err != nil {
return err
}
m.tcpConnDuration, err = meter.Int64Histogram(
"proxy.tcp.connection.duration.ms",
metric.WithUnit("milliseconds"),
metric.WithDescription("Duration of TCP/TLS relay connections"),
)
if err != nil {
return err
}
m.tcpBytesTotal, err = meter.Int64Counter(
"proxy.tcp.bytes.total",
metric.WithUnit("bytes"),
metric.WithDescription("Total bytes transferred through TCP/TLS relay by direction"),
)
if err != nil {
return err
}
m.udpActiveSess, err = meter.Int64UpDownCounter(
"proxy.udp.active_sessions",
metric.WithUnit("1"),
metric.WithDescription("Current number of active UDP relay sessions"),
)
if err != nil {
return err
}
m.udpSessionsTotal, err = meter.Int64Counter(
"proxy.udp.sessions.total",
metric.WithUnit("1"),
metric.WithDescription("Total UDP relay sessions by result and account"),
)
if err != nil {
return err
}
m.udpPacketsTotal, err = meter.Int64Counter(
"proxy.udp.packets.total",
metric.WithUnit("1"),
metric.WithDescription("Total UDP packets relayed by direction"),
)
if err != nil {
return err
}
m.udpBytesTotal, err = meter.Int64Counter(
"proxy.udp.bytes.total",
metric.WithUnit("bytes"),
metric.WithDescription("Total bytes transferred through UDP relay by direction"),
)
return err
}
type responseInterceptor struct {
*responsewriter.PassthroughWriter
status int
size int
}
func (w *responseInterceptor) WriteHeader(status int) {
w.status = status
w.PassthroughWriter.WriteHeader(status)
}
func (w *responseInterceptor) Write(b []byte) (int, error) {
size, err := w.PassthroughWriter.Write(b)
w.size += size
return size, err
}
// Unwrap returns the underlying ResponseWriter so http.ResponseController
// can reach through to the original writer for Hijack/Flush operations.
func (w *responseInterceptor) Unwrap() http.ResponseWriter {
return w.PassthroughWriter
}
// Middleware wraps an HTTP handler with request metrics.
func (m *Metrics) Middleware(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
m.requestsTotal.Add(m.ctx, 1)
m.activeRequests.Add(m.ctx, 1)
interceptor := &responseInterceptor{PassthroughWriter: responsewriter.New(w)}
start := time.Now()
defer func() {
duration := time.Since(start)
m.activeRequests.Add(m.ctx, -1)
m.requestDuration.Record(m.ctx, duration.Milliseconds())
}()
next.ServeHTTP(interceptor, r)
})
}
type roundTripperFunc func(*http.Request) (*http.Response, error)
func (f roundTripperFunc) RoundTrip(r *http.Request) (*http.Response, error) {
return f(r)
}
// RoundTripper wraps an http.RoundTripper with backend duration metrics.
func (m *Metrics) RoundTripper(next http.RoundTripper) http.RoundTripper {
return roundTripperFunc(func(req *http.Request) (*http.Response, error) {
start := time.Now()
res, err := next.RoundTrip(req)
duration := time.Since(start)
m.backendDuration.Record(m.ctx, duration.Milliseconds())
return res, err
})
}
// AddMapping records that a domain mapping was added.
func (m *Metrics) AddMapping(mapping proxy.Mapping) {
m.mappingsMux.Lock()
defer m.mappingsMux.Unlock()
newPathCount := len(mapping.Paths)
oldPathCount, exists := m.mappingPaths[mapping.Host]
if !exists {
m.configuredDomains.Add(m.ctx, 1)
}
pathDelta := newPathCount - oldPathCount
if pathDelta != 0 {
m.totalPaths.Add(m.ctx, int64(pathDelta))
}
m.mappingPaths[mapping.Host] = newPathCount
}
// RemoveMapping records that a domain mapping was removed.
func (m *Metrics) RemoveMapping(mapping proxy.Mapping) {
m.mappingsMux.Lock()
defer m.mappingsMux.Unlock()
oldPathCount, exists := m.mappingPaths[mapping.Host]
if !exists {
return
}
m.configuredDomains.Add(m.ctx, -1)
m.totalPaths.Add(m.ctx, -int64(oldPathCount))
delete(m.mappingPaths, mapping.Host)
}
// RecordCertificateIssuance records the duration of a certificate issuance.
func (m *Metrics) RecordCertificateIssuance(duration time.Duration) {
m.certificateIssueDuration.Record(m.ctx, duration.Milliseconds())
}
// L4ServiceAdded increments the L4 service gauge for the given mode.
func (m *Metrics) L4ServiceAdded(mode types.ServiceMode) {
m.l4Services.Add(m.ctx, 1, metric.WithAttributes(attribute.String("mode", string(mode))))
}
// L4ServiceRemoved decrements the L4 service gauge for the given mode.
func (m *Metrics) L4ServiceRemoved(mode types.ServiceMode) {
m.l4Services.Add(m.ctx, -1, metric.WithAttributes(attribute.String("mode", string(mode))))
}
// TCPRelayStarted records a new TCP relay connection starting.
func (m *Metrics) TCPRelayStarted(accountID types.AccountID) {
acct := attribute.String("account_id", string(accountID))
m.tcpActiveConns.Add(m.ctx, 1, metric.WithAttributes(acct))
m.tcpConnsTotal.Add(m.ctx, 1, metric.WithAttributes(acct, attribute.String("result", "success")))
}
// TCPRelayEnded records a TCP relay connection ending and accumulates bytes and duration.
func (m *Metrics) TCPRelayEnded(accountID types.AccountID, duration time.Duration, srcToDst, dstToSrc int64) {
acct := attribute.String("account_id", string(accountID))
m.tcpActiveConns.Add(m.ctx, -1, metric.WithAttributes(acct))
m.tcpConnDuration.Record(m.ctx, duration.Milliseconds(), metric.WithAttributes(acct))
m.tcpBytesTotal.Add(m.ctx, srcToDst, metric.WithAttributes(attribute.String("direction", "client_to_backend")))
m.tcpBytesTotal.Add(m.ctx, dstToSrc, metric.WithAttributes(attribute.String("direction", "backend_to_client")))
}
// TCPRelayDialError records a dial failure for a TCP relay.
func (m *Metrics) TCPRelayDialError(accountID types.AccountID) {
m.tcpConnsTotal.Add(m.ctx, 1, metric.WithAttributes(
attribute.String("account_id", string(accountID)),
attribute.String("result", "dial_error"),
))
}
// TCPRelayRejected records a rejected TCP relay (semaphore full).
func (m *Metrics) TCPRelayRejected(accountID types.AccountID) {
m.tcpConnsTotal.Add(m.ctx, 1, metric.WithAttributes(
attribute.String("account_id", string(accountID)),
attribute.String("result", "rejected"),
))
}
// UDPSessionStarted records a new UDP session starting.
func (m *Metrics) UDPSessionStarted(accountID types.AccountID) {
acct := attribute.String("account_id", string(accountID))
m.udpActiveSess.Add(m.ctx, 1, metric.WithAttributes(acct))
m.udpSessionsTotal.Add(m.ctx, 1, metric.WithAttributes(acct, attribute.String("result", "success")))
}
// UDPSessionEnded records a UDP session ending.
func (m *Metrics) UDPSessionEnded(accountID types.AccountID) {
m.udpActiveSess.Add(m.ctx, -1, metric.WithAttributes(attribute.String("account_id", string(accountID))))
}
// UDPSessionDialError records a dial failure for a UDP session.
func (m *Metrics) UDPSessionDialError(accountID types.AccountID) {
m.udpSessionsTotal.Add(m.ctx, 1, metric.WithAttributes(
attribute.String("account_id", string(accountID)),
attribute.String("result", "dial_error"),
))
}
// UDPSessionRejected records a rejected UDP session (limit or rate limited).
func (m *Metrics) UDPSessionRejected(accountID types.AccountID) {
m.udpSessionsTotal.Add(m.ctx, 1, metric.WithAttributes(
attribute.String("account_id", string(accountID)),
attribute.String("result", "rejected"),
))
}
// UDPPacketRelayed records a packet relayed in the given direction with its size in bytes.
func (m *Metrics) UDPPacketRelayed(direction types.RelayDirection, bytes int) {
dir := attribute.String("direction", string(direction))
m.udpPacketsTotal.Add(m.ctx, 1, metric.WithAttributes(dir))
m.udpBytesTotal.Add(m.ctx, int64(bytes), metric.WithAttributes(dir))
}
proxy/internal/metrics/metrics_test.go
+83
View File
@@ -0,0 +1,83 @@
package metrics_test
import (
"context"
"net/http"
"net/url"
"reflect"
"testing"
"github.com/google/go-cmp/cmp"
"go.opentelemetry.io/otel/exporters/prometheus"
"go.opentelemetry.io/otel/sdk/metric"
"github.com/netbirdio/netbird/proxy/internal/metrics"
)
type testRoundTripper struct {
response *http.Response
err error
}
func (t *testRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
return t.response, t.err
}
func TestMetrics_RoundTripper(t *testing.T) {
testResponse := http.Response{
StatusCode: http.StatusOK,
Body: http.NoBody,
}
tests := map[string]struct {
roundTripper http.RoundTripper
request *http.Request
response *http.Response
err error
}{
"ok": {
roundTripper: &testRoundTripper{response: &testResponse},
request: &http.Request{Method: "GET", URL: &url.URL{Path: "/foo"}},
response: &testResponse,
},
"nil url": {
roundTripper: &testRoundTripper{response: &testResponse},
request: &http.Request{Method: "GET", URL: nil},
response: &testResponse,
},
"nil response": {
roundTripper: &testRoundTripper{response: nil},
request: &http.Request{Method: "GET", URL: &url.URL{Path: "/foo"}},
},
}
exporter, err := prometheus.New()
if err != nil {
t.Fatalf("create prometheus exporter: %v", err)
}
provider := metric.NewMeterProvider(metric.WithReader(exporter))
pkg := reflect.TypeOf(metrics.Metrics{}).PkgPath()
meter := provider.Meter(pkg)
m, err := metrics.New(context.Background(), meter)
if err != nil {
t.Fatalf("create metrics: %v", err)
}
for name, test := range tests {
t.Run(name, func(t *testing.T) {
rt := m.RoundTripper(test.roundTripper)
res, err := rt.RoundTrip(test.request)
if res != nil && res.Body != nil {
defer res.Body.Close()
}
if diff := cmp.Diff(test.err, err); diff != "" {
t.Errorf("Incorrect error (-want +got):\n%s", diff)
}
if diff := cmp.Diff(test.response, res); diff != "" {
t.Errorf("Incorrect response (-want +got):\n%s", diff)
}
})
}
}
proxy/internal/middleware/bodypolicy.go
+63
View File
@@ -0,0 +1,63 @@
package middleware
import (
"bytes"
"errors"
"fmt"
"io"
"net/http"
"strconv"
"strings"
)
// ErrExpectContinue is returned when a middleware attempts to replace
// the body of a request that advertised Expect: 100-continue.
var ErrExpectContinue = errors.New("body replace rejected: request has Expect: 100-continue")
// ErrOriginalNotDrained is returned when the original body was not
// fully consumed before replacement. This prevents the backend from
// seeing a mix of original bytes and the replacement.
var ErrOriginalNotDrained = errors.New("body replace rejected: original body not drained")
// ErrContentLengthMismatch is returned when the client-advertised
// Content-Length disagrees with the number of bytes actually read from
// the body (short-read).
var ErrContentLengthMismatch = errors.New("body replace rejected: content-length mismatch (short read)")
// ValidateBodyReplace runs the smuggling-prevention rules before a
// body replacement is applied. Callers must pass originalDrained=true
// once they have read r.Body to EOF.
func ValidateBodyReplace(r *http.Request, newBody []byte, originalDrained bool) error {
if r == nil {
return errors.New("body replace rejected: nil request")
}
if strings.EqualFold(r.Header.Get("Expect"), "100-continue") {
return ErrExpectContinue
}
if !originalDrained {
return ErrOriginalNotDrained
}
if cl := r.Header.Get("Content-Length"); cl != "" && r.ContentLength > 0 {
parsed, err := strconv.ParseInt(cl, 10, 64)
if err == nil && parsed != r.ContentLength {
return fmt.Errorf("%w: header=%d actual=%d", ErrContentLengthMismatch, parsed, r.ContentLength)
}
}
return nil
}
// ApplyBodyReplace swaps r.Body for a reader over newBody, recomputes
// Content-Length, and strips Transfer-Encoding and Trailer so no stale
// framing reaches the backend.
func ApplyBodyReplace(r *http.Request, newBody []byte) {
if r == nil {
return
}
r.Body = io.NopCloser(bytes.NewReader(newBody))
r.ContentLength = int64(len(newBody))
r.Header.Set("Content-Length", strconv.Itoa(len(newBody)))
r.Header.Del("Transfer-Encoding")
r.Header.Del("Trailer")
r.TransferEncoding = nil
r.Trailer = nil
}
proxy/internal/middleware/bodytap/request.go
+344
View File
@@ -0,0 +1,344 @@
// Package bodytap owns the framework-side body capture used by the
// middleware chain. Request capture buffers up to N bytes of the
// request body for middleware inspection while replaying the original
// stream to the upstream. Response capture tees up to N bytes off the
// streaming response while every byte continues to flow to the client
// untouched.
//
// The package is the single owner of body access — middlewares never
// read req.Body or hijack the response writer. All inspection happens
// against the buffer surfaced by the tap, so streaming remains
// transparent to the client even when middlewares need access to the
// payload.
package bodytap
import (
"bytes"
"encoding/json"
"errors"
"io"
"net/http"
"strconv"
"strings"
"sync"
)
// MaxRoutingScanBytes bounds how far ScanRoutingFields will read into a
// request body to recover routing fields when the normal capture is
// bypassed for size. Sized to comfortably hold a 1M-token context
// request (whose `model` field a client may place after a multi-MB
// `messages` array) while still capping pathological inputs.
const MaxRoutingScanBytes int64 = 32 << 20
// Request bypass reasons emitted as the `mw.capture.bypass_reason`
// metadata key by the chain when a request body is not surfaced.
const (
BypassUpgradeHeader = "upgrade_header"
BypassConnectionUpgrd = "connection_upgrade"
BypassContentType = "content_type_not_allowed"
BypassBudget = "capture_budget_exhausted"
BypassNoConfig = "no_capture_config"
BypassNoMiddlewares = "no_middlewares"
BypassCapZero = "cap_zero"
BypassContentLengthCap = "content_length_over_cap"
)
// DefaultCaptureBudgetBytes is the default global capture-budget size.
const DefaultCaptureBudgetBytes int64 = 256 << 20
// Config holds per-target body capture limits after clamp validation.
// A zero MaxRequestBytes / MaxResponseBytes disables capture in that
// direction.
type Config struct {
MaxRequestBytes int64
MaxResponseBytes int64
ContentTypes []string
}
// Budget is the global token-bucket semaphore shared across all
// in-flight captures so a single misbehaving target cannot exhaust the
// proxy.
type Budget interface {
Acquire(n int64) bool
Release(n int64)
}
// NewBudget returns a Budget with the given total byte cap. A zero or
// negative total disables the budget check.
func NewBudget(total int64) Budget {
return &budget{total: total}
}
type budget struct {
mu sync.Mutex
used int64
total int64
}
func (b *budget) Acquire(n int64) bool {
if n <= 0 {
return true
}
b.mu.Lock()
defer b.mu.Unlock()
if b.total <= 0 {
return true
}
if b.used+n > b.total {
return false
}
b.used += n
return true
}
func (b *budget) Release(n int64) {
if n <= 0 {
return
}
b.mu.Lock()
defer b.mu.Unlock()
if b.total <= 0 {
return
}
b.used -= n
if b.used < 0 {
b.used = 0
}
}
// CaptureRequest reads up to cfg.MaxRequestBytes from r.Body into a
// buffer suitable for middleware inspection, replacing r.Body with a
// replay reader so the upstream still sees the original bytes. When
// bypass != "" no body is read and r.Body is left untouched. The
// returned release function must be invoked once the request is fully
// processed; it returns the acquired budget tokens to the shared pool.
// release is always non-nil and is safe to defer immediately after the
// call.
func CaptureRequest(r *http.Request, cfg *Config, b Budget) (body []byte, truncated bool, originalSize int64, bypass string, release func(), err error) {
release = func() {}
if r == nil {
return nil, false, 0, BypassNoConfig, release, nil
}
if cfg == nil {
return nil, false, 0, BypassNoConfig, release, nil
}
if cfg.MaxRequestBytes <= 0 {
return nil, false, 0, BypassCapZero, release, nil
}
if r.Header.Get("Upgrade") != "" {
return nil, false, 0, BypassUpgradeHeader, release, nil
}
if strings.EqualFold(r.Header.Get("Connection"), "upgrade") {
return nil, false, 0, BypassConnectionUpgrd, release, nil
}
if !contentTypeAllowed(r.Header.Get("Content-Type"), cfg.ContentTypes) {
return nil, false, 0, BypassContentType, release, nil
}
originalSize = parseContentLength(r.Header.Get("Content-Length"))
if originalSize > cfg.MaxRequestBytes {
return nil, true, originalSize, BypassContentLengthCap, release, nil
}
limit := cfg.MaxRequestBytes
if b != nil && !b.Acquire(limit) {
return nil, false, originalSize, BypassBudget, release, nil
}
if b != nil {
var released sync.Once
release = func() {
released.Do(func() { b.Release(limit) })
}
}
if r.Body == nil || r.Body == http.NoBody {
release()
release = func() {}
return nil, false, originalSize, "", release, nil
}
limited := io.LimitReader(r.Body, limit+1)
buf, readErr := io.ReadAll(limited)
if readErr != nil && !errors.Is(readErr, io.EOF) {
release()
release = func() {}
return nil, false, originalSize, "", release, readErr
}
truncated = int64(len(buf)) > limit
if truncated {
replay := append([]byte(nil), buf...)
viewable := buf[:limit]
r.Body = &replayReadCloser{replay: bytes.NewReader(replay), tail: r.Body}
return viewable, true, originalSize, "", release, nil
}
_ = r.Body.Close()
r.Body = io.NopCloser(bytes.NewReader(buf))
if originalSize <= 0 {
originalSize = int64(len(buf))
}
return buf, false, originalSize, "", release, nil
}
// replayReadCloser replays the captured prefix and then forwards the
// remaining bytes from the original body so the upstream sees the
// full request stream even when capture truncates.
type replayReadCloser struct {
replay *bytes.Reader
tail io.ReadCloser
drained bool
}
func (r *replayReadCloser) Read(p []byte) (int, error) {
if !r.drained {
n, err := r.replay.Read(p)
if n > 0 {
return n, nil
}
if errors.Is(err, io.EOF) {
r.drained = true
} else if err != nil {
return 0, err
}
}
return r.tail.Read(p)
}
func (r *replayReadCloser) Close() error {
return r.tail.Close()
}
// ScanRoutingFields recovers the LLM routing fields ("model" and
// "stream") from a request whose normal capture was bypassed or
// truncated for size. It reads up to maxScan bytes of r.Body to locate
// the top-level keys — clients (e.g. Claude Code) may place `model`
// after a multi-MB `messages` array — then restores r.Body so the
// upstream still receives the full, untouched stream. Only the small
// routing fields are extracted; the prompt is never buffered for
// capture, keeping memory bounded. Returns ok=false when the body isn't
// a JSON object, the model field isn't found within maxScan, or on a
// read error.
func ScanRoutingFields(r *http.Request, maxScan int64) (model string, stream bool, ok bool) {
if r == nil || r.Body == nil || r.Body == http.NoBody || maxScan <= 0 {
return "", false, false
}
limited := io.LimitReader(r.Body, maxScan+1)
buf, readErr := io.ReadAll(limited)
if readErr != nil && !errors.Is(readErr, io.EOF) {
// Mid-stream read error (e.g. client disconnect): restore the bytes
// read so far plus the untouched tail and abort, rather than
// forwarding only the partial prefix as if it were the whole body.
r.Body = &replayReadCloser{replay: bytes.NewReader(append([]byte(nil), buf...)), tail: r.Body}
return "", false, false
}
if int64(len(buf)) > maxScan {
// Body exceeds the scan ceiling: restore the read prefix plus the
// untouched tail so the upstream still gets every byte.
r.Body = &replayReadCloser{replay: bytes.NewReader(append([]byte(nil), buf...)), tail: r.Body}
} else {
_ = r.Body.Close()
r.Body = io.NopCloser(bytes.NewReader(buf))
}
return scanTopLevelModelStream(buf)
}
// scanTopLevelModelStream walks the top level of a JSON object via a
// streaming token reader, extracting the "model" string and "stream"
// bool without materialising large values (each non-target value is
// skipped as a RawMessage). Tolerant of truncation: returns whatever was
// found before a malformed/short tail.
func scanTopLevelModelStream(body []byte) (model string, stream bool, ok bool) {
dec := json.NewDecoder(bytes.NewReader(body))
tok, err := dec.Token()
if err != nil {
return "", false, false
}
if d, isDelim := tok.(json.Delim); !isDelim || d != '{' {
return "", false, false
}
for dec.More() {
keyTok, err := dec.Token()
if err != nil {
return model, stream, ok
}
key, _ := keyTok.(string)
switch key {
case "model":
var v string
if dec.Decode(&v) == nil {
model, ok = v, true
}
case "stream":
var v bool
if dec.Decode(&v) == nil {
stream = v
}
default:
// Skip the value by walking tokens instead of decoding it into
// a json.RawMessage — a multi-MB messages array would otherwise
// be materialised in full just to be discarded.
if err := skipValue(dec); err != nil {
return model, stream, ok
}
}
}
return model, stream, ok
}
// skipValue consumes one JSON value from dec without materialising it.
// Scalars are a single token; objects/arrays are walked to their matching
// close delimiter so nested structures are skipped in bounded memory.
func skipValue(dec *json.Decoder) error {
tok, err := dec.Token()
if err != nil {
return err
}
d, isDelim := tok.(json.Delim)
if !isDelim || (d != '{' && d != '[') {
return nil
}
depth := 1
for depth > 0 {
tok, err := dec.Token()
if err != nil {
return err
}
if d, ok := tok.(json.Delim); ok {
switch d {
case '{', '[':
depth++
case '}', ']':
depth--
}
}
}
return nil
}
func contentTypeAllowed(ct string, allowed []string) bool {
if len(allowed) == 0 {
return false
}
media := ct
if idx := strings.Index(ct, ";"); idx >= 0 {
media = ct[:idx]
}
media = strings.TrimSpace(strings.ToLower(media))
for _, a := range allowed {
if strings.EqualFold(strings.TrimSpace(a), media) {
return true
}
}
return false
}
func parseContentLength(v string) int64 {
if v == "" {
return 0
}
parsed, err := strconv.ParseInt(v, 10, 64)
if err != nil || parsed < 0 {
return 0
}
return parsed
}
proxy/internal/middleware/bodytap/response.go
+189
View File
@@ -0,0 +1,189 @@
package bodytap
import (
"bytes"
"net/http"
"sync"
"github.com/netbirdio/netbird/proxy/internal/responsewriter"
)
// CapturingResponseWriter wraps an http.ResponseWriter, forwards bytes
// immediately to the client, and tees a bounded copy into an internal
// buffer for middleware inspection. Streaming-aware in the sense that
// every byte the upstream emits flows to the client without queuing
// — the tee just sees a bounded prefix. SSE-aware parsing happens in
// the response middleware against the buffered prefix; this writer
// makes no attempt to demux event boundaries.
//
// Flusher and Hijacker are preserved via responsewriter.PassthroughWriter.
type CapturingResponseWriter struct {
*responsewriter.PassthroughWriter
mu sync.Mutex
buf bytes.Buffer
cap int64
status int
statusSet bool
written int64
truncated bool
stopped bool
releaseBuf func()
released sync.Once
bypassed bool
bypassReas string
acquiredCap int64
}
// NewCapturingResponseWriter returns a writer that tees up to maxBytes
// into a capped buffer while forwarding bytes to the underlying writer
// immediately. When budget is non-nil the writer pre-acquires maxBytes
// from it and the returned wrapper must be released by calling
// Release() once the response is fully forwarded. If the budget cannot
// be acquired the writer falls back to forwarding the response
// unmodified, exposes Bypassed()=true with reason BypassBudget, and
// releases nothing.
func NewCapturingResponseWriter(w http.ResponseWriter, maxBytes int64, b Budget) *CapturingResponseWriter {
cw := &CapturingResponseWriter{
PassthroughWriter: responsewriter.New(w),
cap: maxBytes,
status: http.StatusOK,
releaseBuf: func() {},
}
if maxBytes <= 0 {
// Capture disabled: mark stopped so Write never tees and never
// flags truncation (a zero cap means "don't capture", not
// "captured nothing").
cw.stopped = true
return cw
}
if b == nil {
return cw
}
if !b.Acquire(maxBytes) {
cw.bypassed = true
cw.bypassReas = BypassBudget
cw.cap = 0
cw.stopped = true
return cw
}
cw.acquiredCap = maxBytes
cw.releaseBuf = func() { b.Release(maxBytes) }
return cw
}
// Release returns the response capture budget acquired at construction
// back to the shared pool. Idempotent. Safe to call from a defer
// immediately after construction even when the writer ended up
// bypassing the budget.
func (c *CapturingResponseWriter) Release() {
if c == nil {
return
}
c.released.Do(func() {
if c.releaseBuf != nil {
c.releaseBuf()
}
})
}
// Bypassed reports whether the writer fell through to a no-tee
// passthrough because the response capture budget could not be
// acquired.
func (c *CapturingResponseWriter) Bypassed() bool {
if c == nil {
return false
}
c.mu.Lock()
defer c.mu.Unlock()
return c.bypassed
}
// BypassReason returns the bypass code recorded by the budget check.
// Empty when capture proceeded normally.
func (c *CapturingResponseWriter) BypassReason() string {
if c == nil {
return ""
}
c.mu.Lock()
defer c.mu.Unlock()
return c.bypassReas
}
// WriteHeader records the status code and forwards it to the underlying
// writer. Only the first call commits the status — matching HTTP semantics,
// where superfluous WriteHeader calls (and any call after the body has
// started) are ignored — so Status() reflects the code actually sent.
func (c *CapturingResponseWriter) WriteHeader(status int) {
c.mu.Lock()
if c.statusSet {
c.mu.Unlock()
return
}
c.status = status
c.statusSet = true
c.mu.Unlock()
c.PassthroughWriter.WriteHeader(status)
}
// Write forwards p to the underlying writer unmodified and copies up
// to the remaining buffer capacity into the tee buffer.
func (c *CapturingResponseWriter) Write(p []byte) (int, error) {
n, err := c.PassthroughWriter.Write(p)
if n > 0 {
c.mu.Lock()
// The first byte commits the status (implicit 200 if WriteHeader was
// never called); a later WriteHeader must not change Status().
c.statusSet = true
c.written += int64(n)
if !c.stopped {
remaining := c.cap - int64(c.buf.Len())
if remaining <= 0 {
c.truncated = true
c.stopped = true
} else {
take := int64(n)
if take > remaining {
take = remaining
c.truncated = true
c.stopped = true
}
c.buf.Write(p[:take])
}
}
c.mu.Unlock()
}
return n, err
}
// Status returns the captured status code (defaults to 200 when
// WriteHeader has not been called).
func (c *CapturingResponseWriter) Status() int {
c.mu.Lock()
defer c.mu.Unlock()
return c.status
}
// Body returns a copy of the buffered response prefix.
func (c *CapturingResponseWriter) Body() []byte {
c.mu.Lock()
defer c.mu.Unlock()
out := make([]byte, c.buf.Len())
copy(out, c.buf.Bytes())
return out
}
// Truncated reports whether the buffered prefix stopped short of the
// full response stream.
func (c *CapturingResponseWriter) Truncated() bool {
c.mu.Lock()
defer c.mu.Unlock()
return c.truncated
}
// BytesWritten returns the total number of bytes forwarded to the
// underlying writer.
func (c *CapturingResponseWriter) BytesWritten() int64 {
c.mu.Lock()
defer c.mu.Unlock()
return c.written
}
proxy/internal/middleware/bodytap/routing_scan_test.go
+86
View File
@@ -0,0 +1,86 @@
package bodytap
import (
"fmt"
"io"
"net/http/httptest"
"strings"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
// makeBigAnthropicBody builds a request body shaped like Claude Code's:
// a multi-MB "messages" array with the routing fields (model, stream)
// placed AFTER it, which is the ordering that defeats a prefix-only
// capture.
func makeBigAnthropicBody(t *testing.T, model string, stream bool, messagesBytes int) string {
t.Helper()
filler := strings.Repeat("x", messagesBytes)
return fmt.Sprintf(
`{"max_tokens":64000,"messages":[{"role":"user","content":%q}],"model":%q,"stream":%t}`,
filler, model, stream,
)
}
func TestScanRoutingFields_ModelAfterLargeMessages(t *testing.T) {
body := makeBigAnthropicBody(t, "claude-opus-4-8", true, 3<<20) // 3 MiB messages
req := httptest.NewRequest("POST", "https://x/v1/messages", strings.NewReader(body))
model, stream, ok := ScanRoutingFields(req, MaxRoutingScanBytes)
require.True(t, ok, "model must be recovered even when it follows a multi-MB messages array")
assert.Equal(t, "claude-opus-4-8", model, "model field must be extracted")
assert.True(t, stream, "stream field must be extracted")
// Body must be fully restored for the upstream.
got, err := io.ReadAll(req.Body)
require.NoError(t, err)
assert.Equal(t, body, string(got), "the full request body must be replayed to upstream after scanning")
}
func TestScanRoutingFields_SmallBody(t *testing.T) {
body := `{"model":"claude-opus-4-8","stream":false,"messages":[]}`
req := httptest.NewRequest("POST", "https://x/v1/messages", strings.NewReader(body))
model, stream, ok := ScanRoutingFields(req, MaxRoutingScanBytes)
require.True(t, ok)
assert.Equal(t, "claude-opus-4-8", model)
assert.False(t, stream)
got, _ := io.ReadAll(req.Body)
assert.Equal(t, body, string(got), "small bodies must also be restored intact")
}
func TestScanRoutingFields_NoModel(t *testing.T) {
body := `{"stream":true,"messages":[]}`
req := httptest.NewRequest("POST", "https://x/v1/messages", strings.NewReader(body))
_, _, ok := ScanRoutingFields(req, MaxRoutingScanBytes)
assert.False(t, ok, "ok must be false when no model field is present")
got, _ := io.ReadAll(req.Body)
assert.Equal(t, body, string(got), "body must be restored even when model is absent")
}
func TestScanRoutingFields_NotJSON(t *testing.T) {
body := "this is not json at all"
req := httptest.NewRequest("POST", "https://x/v1/messages", strings.NewReader(body))
_, _, ok := ScanRoutingFields(req, MaxRoutingScanBytes)
assert.False(t, ok, "ok must be false for a non-JSON body")
}
func TestScanRoutingFields_ModelBeyondScanCeiling(t *testing.T) {
// model sits after 4 MiB of messages but the scan ceiling is 1 MiB:
// model can't be recovered, yet the full body must still replay.
body := makeBigAnthropicBody(t, "claude-opus-4-8", true, 4<<20)
req := httptest.NewRequest("POST", "https://x/v1/messages", strings.NewReader(body))
_, _, ok := ScanRoutingFields(req, 1<<20)
assert.False(t, ok, "model beyond the scan ceiling is not recoverable")
got, err := io.ReadAll(req.Body)
require.NoError(t, err)
assert.Equal(t, body, string(got), "the full body must still replay to upstream even when the scan gives up")
}
proxy/internal/middleware/builtin/agentnetwork_chain_integration_test.go
+318
View File
@@ -0,0 +1,318 @@
package builtin_test
import (
"context"
"net"
"runtime"
"strconv"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"google.golang.org/grpc"
"google.golang.org/grpc/credentials/insecure"
"google.golang.org/grpc/test/bufconn"
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork"
agentNetworkTypes "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
mgmtgrpc "github.com/netbirdio/netbird/management/internals/shared/grpc"
"github.com/netbirdio/netbird/management/server/store"
nbtypes "github.com/netbirdio/netbird/management/server/types"
"github.com/netbirdio/netbird/proxy/internal/middleware"
"github.com/netbirdio/netbird/proxy/internal/middleware/builtin/llm_limit_check"
"github.com/netbirdio/netbird/proxy/internal/middleware/builtin/llm_limit_record"
"github.com/netbirdio/netbird/shared/management/proto"
)
// chainIntegrationFixture wires the BOTH new agent-network
// middlewares against a live in-process management stack: real
// sqlite store + real Manager + real gRPC server. The proxy chain
// framework itself isn't constructed (its dispatcher / accumulator /
// metadata gate are tested separately); we exercise the middleware
// pair as the proxy runtime would, by invoking each with a crafted
// Input and asserting the wire path between them.
//
// This is the regression cover for item 16 in the design review:
// real LLM request → cost stamped → consumption row in the table.
type chainIntegrationFixture struct {
store store.Store
manager agentnetwork.Manager
gatecase *llm_limit_check.Middleware
recorder *llm_limit_record.Middleware
}
func newChainIntegration(t *testing.T) *chainIntegrationFixture {
t.Helper()
if runtime.GOOS == "windows" {
t.Skip("sqlite store not properly supported on Windows yet")
}
t.Setenv("NETBIRD_STORE_ENGINE", string(nbtypes.SqliteStoreEngine))
st, cleanUp, err := store.NewTestStoreFromSQL(context.Background(), "", t.TempDir())
require.NoError(t, err)
t.Cleanup(cleanUp)
manager := agentnetwork.NewManager(st, nil, nil, nil)
server := &mgmtgrpc.ProxyServiceServer{}
server.SetAgentNetworkLimitsService(manager)
const bufSize = 1024 * 1024
lis := bufconn.Listen(bufSize)
srv := grpc.NewServer()
proto.RegisterProxyServiceServer(srv, server)
go func() { _ = srv.Serve(lis) }()
t.Cleanup(srv.Stop)
conn, err := grpc.NewClient("passthrough:///bufnet",
grpc.WithContextDialer(func(_ context.Context, _ string) (net.Conn, error) { return lis.Dial() }),
grpc.WithTransportCredentials(insecure.NewCredentials()),
)
require.NoError(t, err)
t.Cleanup(func() { _ = conn.Close() })
mgmtClient := proto.NewProxyServiceClient(conn)
return &chainIntegrationFixture{
store: st,
manager: manager,
gatecase: llm_limit_check.New(mgmtClient, nil),
recorder: llm_limit_record.New(mgmtClient, nil),
}
}
// chainInput builds a middleware Input that mirrors what the proxy
// framework would synthesise for a tunnel-peer LLM request. The
// gate consumes the resolved provider id from upstream metadata
// (set by llm_router); the recorder consumes the attribution
// metadata stamped by the gate plus tokens / cost from
// llm_response_parser + cost_meter.
func chainInput(account, user, group, providerID string, requestMeta []middleware.KV) *middleware.Input {
_ = providerID // packed into requestMeta by the caller as KeyLLMResolvedProviderID
return &middleware.Input{
AccountID: account,
UserID: user,
UserGroups: []string{group},
Metadata: requestMeta,
}
}
// chainCapPolicy builds a tight token-cap policy fixture for the
// chain integration tests. Inlined here (rather than imported) because
// the equivalent helper in the management gRPC package is unexported
// and this is a different package boundary.
func chainCapPolicy(id, account string, sourceGroups []string, providerID string, tokenCap, windowSec int64) *agentNetworkTypes.Policy {
return &agentNetworkTypes.Policy{
ID: id,
AccountID: account,
Enabled: true,
Name: id,
SourceGroups: sourceGroups,
DestinationProviderIDs: []string{providerID},
Limits: agentNetworkTypes.PolicyLimits{
TokenLimit: agentNetworkTypes.PolicyTokenLimit{
Enabled: true,
GroupCap: tokenCap,
WindowSeconds: windowSec,
},
},
CreatedAt: time.Now().UTC(),
UpdatedAt: time.Now().UTC(),
}
}
// TestChain_AllowPath_StampsAttributionAndRecordsCounter walks the
// full happy path: gate calls CheckLLMPolicyLimits → stamps
// attribution metadata → recorder reads metadata + tokens / cost →
// calls RecordLLMUsage → counters land in sqlite. Asserting on the
// store at the end proves every leg of the wire works together,
// not just each leg in isolation (which the unit tests already cover).
func TestChain_AllowPath_StampsAttributionAndRecordsCounter(t *testing.T) {
f := newChainIntegration(t)
const account = "acc-1"
const user = "user-bob"
const group = "grp-engineers"
const provider = "prov-1"
// Seed a policy with token + budget caps; both halves carry
// real ceilings so the request stays within headroom.
require.NoError(t, f.store.SaveAgentNetworkPolicy(context.Background(),
chainCapPolicy("pol-1", account, []string{group}, provider, 10_000, 86_400)))
// ── Stage 1 — gate: pre-flight check ──────────────────────
gateIn := chainInput(account, user, group, provider, []middleware.KV{
{Key: middleware.KeyLLMResolvedProviderID, Value: provider},
{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
})
gateOut, err := f.gatecase.Invoke(context.Background(), gateIn)
require.NoError(t, err)
require.Equal(t, middleware.DecisionAllow, gateOut.Decision, "fresh policy must allow")
// Verify attribution metadata was stamped — the recorder
// depends on these keys.
metaMap := map[string]string{}
for _, kv := range gateOut.Metadata {
metaMap[kv.Key] = kv.Value
}
assert.Equal(t, "pol-1", metaMap[middleware.KeyLLMSelectedPolicyID])
assert.Equal(t, group, metaMap[middleware.KeyLLMAttributionGroupID])
assert.Equal(t, "86400", metaMap[middleware.KeyLLMAttributionWindowS])
// ── Stage 2 — recorder: post-flight write ─────────────────
// Build the response-leg Input the framework would synthesise
// for the recorder: gate's emitted attribution metadata + the
// tokens / cost stamped by llm_response_parser + cost_meter.
const tokensIn = int64(123)
const tokensOut = int64(45)
const costUSD = 0.0042
recordIn := chainInput(account, user, group, provider, append([]middleware.KV{},
gateOut.Metadata...))
recordIn.Metadata = append(recordIn.Metadata,
middleware.KV{Key: middleware.KeyLLMInputTokens, Value: strconv.FormatInt(tokensIn, 10)},
middleware.KV{Key: middleware.KeyLLMOutputTokens, Value: strconv.FormatInt(tokensOut, 10)},
middleware.KV{Key: middleware.KeyCostUSDTotal, Value: strconv.FormatFloat(costUSD, 'f', 6, 64)},
)
recordOut, err := f.recorder.Invoke(context.Background(), recordIn)
require.NoError(t, err)
assert.Equal(t, middleware.DecisionAllow, recordOut.Decision, "recorder always allows; its only side effect is the counter write")
// ── Stage 3 — assert state in sqlite ──────────────────────
windowStart := agentNetworkTypes.WindowStart(time.Now(), 86_400)
userRow, err := f.store.GetAgentNetworkConsumption(
context.Background(), store.LockingStrengthNone, account,
agentNetworkTypes.DimensionUser, user, int64(86_400), windowStart,
)
require.NoError(t, err)
assert.Equal(t, tokensIn, userRow.TokensInput, "user counter must hold the input tokens the recorder posted")
assert.Equal(t, tokensOut, userRow.TokensOutput)
assert.InDelta(t, costUSD, userRow.CostUSD, 1e-6)
groupRow, err := f.store.GetAgentNetworkConsumption(
context.Background(), store.LockingStrengthNone, account,
agentNetworkTypes.DimensionGroup, group, int64(86_400), windowStart,
)
require.NoError(t, err)
assert.Equal(t, tokensIn, groupRow.TokensInput, "group counter mirrors the user counter — single Record posts both dims")
}
// TestChain_DenyPath_GateRejectsAndNoConsumptionWritten covers the
// negative side: when the gate denies, the recorder is never
// invoked (the proxy framework short-circuits on Decision=Deny).
// We assert no consumption row materialises after the gate-deny
// path, even though the test technically calls the recorder
// afterwards — the recorder must skip on missing attribution
// metadata so the framework's short-circuit isn't load-bearing for
// data integrity.
func TestChain_DenyPath_GateRejectsAndNoConsumptionWritten(t *testing.T) {
f := newChainIntegration(t)
const account = "acc-1"
const user = "user-bob"
const group = "grp-tight"
const provider = "prov-1"
policy := chainCapPolicy("pol-tight", account, []string{group}, provider, 100, 86_400)
require.NoError(t, f.store.SaveAgentNetworkPolicy(context.Background(), policy))
// Pre-burn the counter to the cap so the gate denies.
require.NoError(t, f.store.IncrementAgentNetworkConsumption(
context.Background(), account,
agentNetworkTypes.DimensionGroup, group, int64(86_400),
agentNetworkTypes.WindowStart(time.Now(), 86_400),
100, 0, 0,
))
gateOut, err := f.gatecase.Invoke(context.Background(), chainInput(account, user, group, provider,
[]middleware.KV{
{Key: middleware.KeyLLMResolvedProviderID, Value: provider},
{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
},
))
require.NoError(t, err)
assert.Equal(t, middleware.DecisionDeny, gateOut.Decision, "policy at-cap must deny on the gate")
require.NotNil(t, gateOut.DenyReason)
assert.Equal(t, "llm_policy.token_cap_exceeded", gateOut.DenyReason.Code)
// On deny, the gate emits no attribution metadata. If the
// proxy framework still invokes the recorder (defense in
// depth), the recorder's "no attribution window = skip" guard
// prevents a phantom counter increment.
recordOut, err := f.recorder.Invoke(context.Background(), chainInput(account, user, group, provider,
gateOut.Metadata, // no llm.attribution_window_seconds stamped
))
require.NoError(t, err)
assert.Equal(t, middleware.DecisionAllow, recordOut.Decision)
// The pre-burned 100 tokens are the only counter movement —
// the recorder must NOT have added a fresh row for the user
// dimension on this denied request.
windowStart := agentNetworkTypes.WindowStart(time.Now(), 86_400)
userRow, err := f.store.GetAgentNetworkConsumption(
context.Background(), store.LockingStrengthNone, account,
agentNetworkTypes.DimensionUser, user, int64(86_400), windowStart,
)
require.NoError(t, err)
assert.Zero(t, userRow.TokensInput, "user dimension must not gain tokens from a denied request — recorder skip is the safety net")
}
// TestChain_CapExhaustTransition exercises the allow→deny boundary
// the operator cares most about: a request just under cap allows
// AND records, the next request post-record at-cap denies. This is
// the same lifecycle 50-grpc-allow-record-deny.sh runs in bash, but
// against the actual middleware pair rather than the smoke binary
// driving the gRPC RPCs directly.
func TestChain_CapExhaustTransition(t *testing.T) {
f := newChainIntegration(t)
const account = "acc-1"
const user = "user-alice"
const group = "grp-cap-edge"
const provider = "prov-1"
const tightCap = int64(100)
require.NoError(t, f.store.SaveAgentNetworkPolicy(context.Background(),
chainCapPolicy("pol-edge", account, []string{group}, provider, tightCap, 86_400)))
// Pre-burn 99 tokens so we're at the very edge.
require.NoError(t, f.store.IncrementAgentNetworkConsumption(
context.Background(), account,
agentNetworkTypes.DimensionGroup, group, int64(86_400),
agentNetworkTypes.WindowStart(time.Now(), 86_400),
99, 0, 0,
))
// Gate at 99/100 — must allow (one token of headroom).
gateOut, err := f.gatecase.Invoke(context.Background(), chainInput(account, user, group, provider,
[]middleware.KV{
{Key: middleware.KeyLLMResolvedProviderID, Value: provider},
{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
},
))
require.NoError(t, err)
require.Equal(t, middleware.DecisionAllow, gateOut.Decision, "99/100 must allow — one token of headroom")
// Record one more input token — pushes us to 100/100.
recordIn := chainInput(account, user, group, provider, append([]middleware.KV{},
gateOut.Metadata...))
recordIn.Metadata = append(recordIn.Metadata,
middleware.KV{Key: middleware.KeyLLMInputTokens, Value: "1"},
middleware.KV{Key: middleware.KeyLLMOutputTokens, Value: "0"},
middleware.KV{Key: middleware.KeyCostUSDTotal, Value: "0.000001"},
)
_, err = f.recorder.Invoke(context.Background(), recordIn)
require.NoError(t, err)
// Next gate call must deny — counter is exactly at cap.
gateOut2, err := f.gatecase.Invoke(context.Background(), chainInput(account, user, group, provider,
[]middleware.KV{
{Key: middleware.KeyLLMResolvedProviderID, Value: provider},
{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
},
))
require.NoError(t, err)
assert.Equal(t, middleware.DecisionDeny, gateOut2.Decision,
"once recorder pushed the group counter to 100/100, the next gate call must deny — allow→deny transition is the operator-visible product semantic")
require.NotNil(t, gateOut2.DenyReason)
assert.Equal(t, "llm_policy.token_cap_exceeded", gateOut2.DenyReason.Code)
}
proxy/internal/middleware/builtin/all_test.go
+40
View File
@@ -0,0 +1,40 @@
package builtin_test
import (
"sort"
"testing"
"github.com/stretchr/testify/assert"
mwbuiltin "github.com/netbirdio/netbird/proxy/internal/middleware/builtin"
_ "github.com/netbirdio/netbird/proxy/internal/middleware/builtin/cost_meter"
_ "github.com/netbirdio/netbird/proxy/internal/middleware/builtin/llm_guardrail"
_ "github.com/netbirdio/netbird/proxy/internal/middleware/builtin/llm_identity_inject"
_ "github.com/netbirdio/netbird/proxy/internal/middleware/builtin/llm_limit_check"
_ "github.com/netbirdio/netbird/proxy/internal/middleware/builtin/llm_limit_record"
_ "github.com/netbirdio/netbird/proxy/internal/middleware/builtin/llm_request_parser"
_ "github.com/netbirdio/netbird/proxy/internal/middleware/builtin/llm_response_parser"
_ "github.com/netbirdio/netbird/proxy/internal/middleware/builtin/llm_router"
)
// TestDefaultRegistry_BuiltinIDs locks the set of middleware IDs that
// the default builtin registry exposes once every sub-package's init()
// has run. The list is the source of truth wired by the synthesiser
// in management; adding a new built-in middleware should consciously
// extend this list.
func TestDefaultRegistry_BuiltinIDs(t *testing.T) {
got := mwbuiltin.DefaultRegistry().IDs()
sort.Strings(got)
want := []string{
"cost_meter",
"llm_guardrail",
"llm_identity_inject",
"llm_limit_check",
"llm_limit_record",
"llm_request_parser",
"llm_response_parser",
"llm_router",
}
assert.Equal(t, want, got, "default registry must expose every built-in middleware after anonymous imports")
}
proxy/internal/middleware/builtin/builtin.go
+93
View File
@@ -0,0 +1,93 @@
// Package builtin holds the package-level middleware registry that
// concrete middleware packages register themselves into via init().
// Server boot anonymous-imports each middleware sub-package; the
// resolver attached to the middleware Manager pulls factories out of
// this registry.
package builtin
import (
"context"
"sync"
log "github.com/sirupsen/logrus"
"go.opentelemetry.io/otel/metric"
"google.golang.org/grpc"
"github.com/netbirdio/netbird/proxy/internal/middleware"
"github.com/netbirdio/netbird/shared/management/proto"
)
// MgmtClient is the narrow slice of proto.ProxyServiceClient that
// builtin middlewares may use during request / response handling.
// Only the agent-network limit pair (llm_limit_check + llm_limit_record)
// uses this today; declaring the surface here keeps the dependency
// explicit at boot time.
//
// proto.ProxyServiceClient already satisfies this interface so server
// boot just forwards its existing client.
type MgmtClient interface {
CheckLLMPolicyLimits(ctx context.Context, in *proto.CheckLLMPolicyLimitsRequest, opts ...grpc.CallOption) (*proto.CheckLLMPolicyLimitsResponse, error)
RecordLLMUsage(ctx context.Context, in *proto.RecordLLMUsageRequest, opts ...grpc.CallOption) (*proto.RecordLLMUsageResponse, error)
}
// defaultRegistry is the package-level registry that concrete builtin
// middlewares register themselves into via init().
var defaultRegistry = middleware.NewRegistry()
// FactoryContext is the per-process bag that concrete factories may
// consult during construction. It carries the proxy-lifetime context,
// the data directory used for static config files (pricing tables,
// allowlists), the OTel meter, and the proxy logger.
//
// Configure must be called once at boot before any chain build calls
// Resolve. Calling it twice overwrites the prior value; tests may rely
// on this to reset state.
type FactoryContext struct {
Context context.Context
DataDir string
Meter metric.Meter
Logger *log.Logger
MgmtClient MgmtClient
}
var (
ctxStore FactoryContext
ctxMu sync.RWMutex
)
// Configure stores the per-process FactoryContext. Concrete factories
// reach for it via Context(). mgmt may be nil on tests / standalone
// builds with no management server; consumers must guard.
func Configure(ctx context.Context, dataDir string, meter metric.Meter, logger *log.Logger, mgmt MgmtClient) {
ctxMu.Lock()
defer ctxMu.Unlock()
ctxStore = FactoryContext{
Context: ctx,
DataDir: dataDir,
Meter: meter,
Logger: logger,
MgmtClient: mgmt,
}
}
// Context returns the stored FactoryContext. Returns a zero value when
// Configure was never called; consumers must guard against nil
// Context/Meter/Logger if they care.
func Context() FactoryContext {
ctxMu.RLock()
defer ctxMu.RUnlock()
return ctxStore
}
// Register adds a factory to the default registry. Called from init()
// blocks of concrete middleware packages. Panics on collision so
// duplicate IDs surface at startup.
func Register(f middleware.Factory) {
defaultRegistry.MustRegister(f)
}
// DefaultRegistry returns the shared registry. The proxy server
// constructs the Resolver from it at boot.
func DefaultRegistry() *middleware.Registry {
return defaultRegistry
}
proxy/internal/middleware/builtin/cost_calculation_matrix_test.go
+329
View File
@@ -0,0 +1,329 @@
package builtin_test
import (
"bytes"
"context"
"encoding/base64"
"encoding/json"
"strconv"
"testing"
"github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/netbirdio/netbird/proxy/internal/middleware"
"github.com/netbirdio/netbird/proxy/internal/middleware/builtin"
"github.com/netbirdio/netbird/proxy/internal/middleware/builtin/cost_meter"
"github.com/netbirdio/netbird/proxy/internal/middleware/builtin/llm_request_parser"
"github.com/netbirdio/netbird/proxy/internal/middleware/builtin/llm_response_parser"
)
// Drives the real pipeline (llm_request_parser → llm_response_parser → cost_meter) on the embedded default pricing
// table and asserts exact USD amounts hardcoded from the vendors' published prices, including the cache split.
func TestCostCalculation_ProviderMatrix(t *testing.T) {
// Empty data dir → embedded defaults, like a proxy with no pricing override.
builtin.Configure(context.Background(), t.TempDir(), nil, nil, nil)
reqMW, err := llm_request_parser.Factory{}.New(nil)
require.NoError(t, err, "build llm_request_parser")
respMW, err := llm_response_parser.Factory{}.New(nil)
require.NoError(t, err, "build llm_response_parser")
costMW, err := cost_meter.Factory{}.New(nil)
require.NoError(t, err, "build cost_meter")
t.Cleanup(func() { _ = costMW.Close() })
const jsonCT = "application/json"
const sseCT = "text/event-stream"
const awsCT = "application/vnd.amazon.eventstream"
cases := []struct {
name string
url string
reqBody []byte
respCT string
respBody []byte
wantProvider string
wantModel string
wantCost float64 // exact expected USD; ignored when wantSkip is set
wantCacheCost float64 // expected cost.usd_cache portion of wantCost
wantSkip string // expected cost.skipped reason, "" when priced
}{
{
// gpt-4o-mini $0.15/$0.60 per MTok: 1000×0.15/1M + 500×0.60/1M.
name: "openai chat completions",
url: "https://api.openai.com/v1/chat/completions",
reqBody: []byte(`{"model":"gpt-4o-mini","messages":[{"role":"user","content":"hi"}]}`),
respCT: jsonCT,
respBody: []byte(`{"choices":[{"message":{"content":"pong"}}],"usage":{"prompt_tokens":1000,"completion_tokens":500,"total_tokens":1500}}`),
wantProvider: "openai",
wantModel: "gpt-4o-mini",
wantCost: 0.00045,
},
{
// OpenAI cached tokens are a SUBSET of prompt_tokens at a discount; gpt-4o $2.50/$10 per MTok, cached $1.25/M:
// 250×2.5/1M + 750×1.25/1M + 500×10/1M.
name: "openai cached subset discount",
url: "https://api.openai.com/v1/chat/completions",
reqBody: []byte(`{"model":"gpt-4o","messages":[{"role":"user","content":"hi"}]}`),
respCT: jsonCT,
respBody: []byte(`{"usage":{"prompt_tokens":1000,"completion_tokens":500,"prompt_tokens_details":{"cached_tokens":750}}}`),
wantProvider: "openai",
wantModel: "gpt-4o",
wantCost: 0.0065625,
wantCacheCost: 0.0009375,
},
{
// OpenAI streaming: usage rides the final SSE frame.
name: "openai chat SSE stream",
url: "https://api.openai.com/v1/chat/completions",
reqBody: []byte(`{"model":"gpt-4o-mini","stream":true,"messages":[{"role":"user","content":"hi"}]}`),
respCT: sseCT,
respBody: sseBody(`{"choices":[{"delta":{"content":"po"}}]}`, `{"choices":[{"delta":{"content":"ng"}}]}`, `{"choices":[],"usage":{"prompt_tokens":1000,"completion_tokens":500}}`, "[DONE]"),
wantProvider: "openai",
wantModel: "gpt-4o-mini",
wantCost: 0.00045,
},
{
// Mistral speaks the OpenAI shape: mistral-large-latest $0.50/$1.50 per MTok.
name: "mistral via openai shape",
url: "https://api.mistral.ai/v1/chat/completions",
reqBody: []byte(`{"model":"mistral-large-latest","messages":[{"role":"user","content":"hi"}]}`),
respCT: jsonCT,
respBody: []byte(`{"usage":{"prompt_tokens":1000,"completion_tokens":1000}}`),
wantProvider: "openai",
wantModel: "mistral-large-latest",
wantCost: 0.002,
},
{
// The field report, minus caching: Bedrock Sonnet 4.6 $3/$15 per MTok, 3×3/1M + 1514×15/1M = $0.022719.
// Also covers inference-profile normalization of the region-prefixed versioned id in the URL.
name: "bedrock invoke — reported scenario, no cache",
url: "https://bedrock-runtime.eu-central-1.amazonaws.com/model/global.anthropic.claude-sonnet-4-6-20260115-v1:0/invoke",
reqBody: []byte(`{"messages":[{"role":"user","content":"hi"}]}`),
respCT: jsonCT,
respBody: []byte(`{"content":[{"type":"text","text":"pong"}],"usage":{"input_tokens":3,"output_tokens":1514}}`),
wantProvider: "bedrock",
wantModel: "anthropic.claude-sonnet-4-6",
wantCost: 0.022719,
},
{
// The field report as observed: the FIRST call of a session also wrote a 30,528-token prompt cache at
// 1.25× input ($3.75/M): 0.022719 + 30528×3.75/1M = $0.137199 — the reported $0.1372.
name: "bedrock invoke — reported scenario with cache write",
url: "https://bedrock-runtime.eu-central-1.amazonaws.com/model/global.anthropic.claude-sonnet-4-6-20260115-v1:0/invoke",
reqBody: []byte(`{"messages":[{"role":"user","content":"hi"}]}`),
respCT: jsonCT,
respBody: []byte(`{"usage":{"input_tokens":3,"output_tokens":1514,"cache_creation_input_tokens":30528,"cache_read_input_tokens":0}}`),
wantProvider: "bedrock",
wantModel: "anthropic.claude-sonnet-4-6",
wantCost: 0.137199,
wantCacheCost: 0.11448,
},
{
// Same numbers over the InvokeModel event-stream: message_start carries input + cache, message_delta the output.
name: "bedrock invoke stream with cache write",
url: "https://bedrock-runtime.eu-central-1.amazonaws.com/model/global.anthropic.claude-sonnet-4-6-20260115-v1:0/invoke-with-response-stream",
reqBody: []byte(`{"messages":[{"role":"user","content":"hi"}]}`),
respCT: awsCT,
respBody: bedrockInvokeStream(t, `{"type":"message_start","message":{"usage":{"input_tokens":3,"output_tokens":1,"cache_creation_input_tokens":30528}}}`, `{"type":"content_block_delta","delta":{"type":"text_delta","text":"pong"}}`, `{"type":"message_delta","usage":{"output_tokens":1514}}`),
wantProvider: "bedrock",
wantModel: "anthropic.claude-sonnet-4-6",
wantCost: 0.137199,
wantCacheCost: 0.11448,
},
{
// Converse camelCase usage incl. cache buckets. Haiku 4.5 $1/$5 per MTok, read $0.10/M, write $1.25/M:
// 50×1/1M + 100×5/1M + 2000×0.1/1M + 1000×1.25/1M = $0.002.
name: "bedrock converse with cache buckets",
url: "https://bedrock-runtime.eu-central-1.amazonaws.com/model/eu.anthropic.claude-haiku-4-5-20251001-v1:0/converse",
reqBody: []byte(`{"messages":[{"role":"user","content":[{"text":"hi"}]}]}`),
respCT: jsonCT,
respBody: []byte(`{"output":{"message":{"content":[{"text":"pong"}]}},"usage":{"inputTokens":50,"outputTokens":100,"totalTokens":3150,"cacheReadInputTokens":2000,"cacheWriteInputTokens":1000}}`),
wantProvider: "bedrock",
wantModel: "anthropic.claude-haiku-4-5",
wantCost: 0.002,
wantCacheCost: 0.00145,
},
{
// Same numbers over converse-stream: usage rides the trailing metadata frame.
name: "bedrock converse stream with cache buckets",
url: "https://bedrock-runtime.eu-central-1.amazonaws.com/model/eu.anthropic.claude-haiku-4-5-20251001-v1:0/converse-stream",
reqBody: []byte(`{"messages":[{"role":"user","content":[{"text":"hi"}]}]}`),
respCT: awsCT,
respBody: bedrockConverseStream(t,
`{"delta":{"text":"pong"}}`,
`{"usage":{"inputTokens":50,"outputTokens":100,"totalTokens":3150,"cacheReadInputTokens":2000,"cacheWriteInputTokens":1000}}`,
),
wantProvider: "bedrock",
wantModel: "anthropic.claude-haiku-4-5",
wantCost: 0.002,
wantCacheCost: 0.00145,
},
{
// First-party Anthropic, additive cache buckets. Sonnet 4.6:
// 256×3/1M + 200×15/1M + 768×0.3/1M + 512×3.75/1M.
name: "anthropic messages with cache buckets",
url: "https://api.anthropic.com/v1/messages",
reqBody: []byte(`{"model":"claude-sonnet-4-6","messages":[{"role":"user","content":"hi"}]}`),
respCT: jsonCT,
respBody: []byte(`{"content":[{"type":"text","text":"pong"}],"usage":{"input_tokens":256,"output_tokens":200,"cache_read_input_tokens":768,"cache_creation_input_tokens":512}}`),
wantProvider: "anthropic",
wantModel: "claude-sonnet-4-6",
wantCost: 0.0059184,
wantCacheCost: 0.0021504,
},
{
// Anthropic SSE: input from message_start, output from message_delta. Haiku 4.5: 1000×1/1M + 2000×5/1M.
name: "anthropic SSE stream",
url: "https://api.anthropic.com/v1/messages",
reqBody: []byte(`{"model":"claude-haiku-4-5","stream":true,"messages":[{"role":"user","content":"hi"}]}`),
respCT: sseCT,
respBody: sseBody(`{"type":"message_start","message":{"usage":{"input_tokens":1000,"output_tokens":2}}}`, `{"type":"content_block_delta","delta":{"type":"text_delta","text":"pong"}}`, `{"type":"message_delta","usage":{"output_tokens":2000}}`, `{"type":"message_stop"}`),
wantProvider: "anthropic",
wantModel: "claude-haiku-4-5",
wantCost: 0.011,
},
{
// Kimi's Anthropic-compatible endpoint: kimi-k3 $3/$15 per MTok under the anthropic table.
name: "kimi anthropic shape",
url: "https://api.moonshot.ai/anthropic/v1/messages",
reqBody: []byte(`{"model":"kimi-k3","messages":[{"role":"user","content":"hi"}]}`),
respCT: jsonCT,
respBody: []byte(`{"usage":{"input_tokens":1000,"output_tokens":1000}}`),
wantProvider: "anthropic",
wantModel: "kimi-k3",
wantCost: 0.018,
},
{
// Vertex path-routed model with "@version" stripped; Anthropic-on-Vertex priced under the anthropic table.
name: "vertex anthropic path-routed",
url: "https://aiplatform.googleapis.com/v1/projects/p/locations/global/publishers/anthropic/models/claude-sonnet-4-6@20260115:rawPredict",
reqBody: []byte(`{"anthropic_version":"vertex-2023-10-16","messages":[{"role":"user","content":"hi"}]}`),
respCT: jsonCT,
respBody: []byte(`{"usage":{"input_tokens":200,"output_tokens":100}}`),
wantProvider: "anthropic",
wantModel: "claude-sonnet-4-6",
wantCost: 0.0021,
},
{
// Gateway-prefixed model ids are not in the pricing table: the meter must SKIP, never guess a rate.
name: "gateway-prefixed model skips pricing",
url: "https://gateway.example.com/v1/chat/completions",
reqBody: []byte(`{"model":"openai/gpt-4o-mini","messages":[{"role":"user","content":"hi"}]}`),
respCT: jsonCT,
respBody: []byte(`{"usage":{"prompt_tokens":1000,"completion_tokens":500}}`),
wantProvider: "openai",
wantModel: "openai/gpt-4o-mini",
wantSkip: "unknown_model",
},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
in := &middleware.Input{
Method: "POST",
URL: tc.url,
Headers: []middleware.KV{{Key: "Content-Type", Value: "application/json"}},
Body: tc.reqBody,
}
reqOut, err := reqMW.Invoke(context.Background(), in)
require.NoError(t, err, "request parser")
in.Metadata = append(in.Metadata, reqOut.Metadata...)
require.Equal(t, tc.wantProvider, metaKV(in.Metadata, middleware.KeyLLMProvider), "detected provider")
require.Equal(t, tc.wantModel, metaKV(in.Metadata, middleware.KeyLLMModel), "detected (normalized) model")
in.Status = 200
in.RespHeaders = []middleware.KV{{Key: "Content-Type", Value: tc.respCT}}
in.RespBody = tc.respBody
respOut, err := respMW.Invoke(context.Background(), in)
require.NoError(t, err, "response parser")
in.Metadata = append(in.Metadata, respOut.Metadata...)
costOut, err := costMW.Invoke(context.Background(), in)
require.NoError(t, err, "cost meter")
if tc.wantSkip != "" {
assert.Equal(t, tc.wantSkip, metaKV(costOut.Metadata, middleware.KeyCostSkipped), "expected cost skip reason")
assert.Empty(t, metaKV(costOut.Metadata, middleware.KeyCostUSDTotal), "no cost may be emitted on skip")
return
}
raw := metaKV(costOut.Metadata, middleware.KeyCostUSDTotal)
require.NotEmpty(t, raw, "cost.usd_total must be emitted; skip=%q", metaKV(costOut.Metadata, middleware.KeyCostSkipped))
got, err := strconv.ParseFloat(raw, 64)
require.NoError(t, err, "cost must be a float")
// cost.usd_total is rendered with %.6f: allow half of the last printed digit on top of float error.
assert.InDelta(t, tc.wantCost, got, 5.1e-7, "USD cost for %s", tc.name)
rawCache := metaKV(costOut.Metadata, middleware.KeyCostUSDCache)
require.NotEmpty(t, rawCache, "cost.usd_cache must be emitted next to cost.usd_total")
gotCache, err := strconv.ParseFloat(rawCache, 64)
require.NoError(t, err, "cache cost must be a float")
assert.InDelta(t, tc.wantCacheCost, gotCache, 5.1e-7, "cache USD cost for %s", tc.name)
})
}
}
// metaKV returns the value for key in kvs, or "" when absent.
func metaKV(kvs []middleware.KV, key string) string {
for _, kv := range kvs {
if kv.Key == key {
return kv.Value
}
}
return ""
}
// sseBody renders data frames as a text/event-stream body.
func sseBody(frames ...string) []byte {
var b bytes.Buffer
for _, f := range frames {
b.WriteString("data: ")
b.WriteString(f)
b.WriteString("\n\n")
}
return b.Bytes()
}
// awsFrame encodes one AWS event-stream frame with the given :event-type.
func awsFrame(t *testing.T, eventType string, payload []byte) []byte {
t.Helper()
var buf bytes.Buffer
enc := eventstream.NewEncoder()
require.NoError(t, enc.Encode(&buf, eventstream.Message{
Headers: eventstream.Headers{{Name: ":event-type", Value: eventstream.StringValue(eventType)}},
Payload: payload,
}), "encode event-stream frame")
return buf.Bytes()
}
// bedrockInvokeStream builds an invoke-with-response-stream body: each "chunk" frame wraps a base64 Anthropic event.
func bedrockInvokeStream(t *testing.T, events ...string) []byte {
t.Helper()
var body bytes.Buffer
for _, ev := range events {
wrap, err := json.Marshal(map[string]string{"bytes": base64.StdEncoding.EncodeToString([]byte(ev))})
require.NoError(t, err)
body.Write(awsFrame(t, "chunk", wrap))
}
return body.Bytes()
}
// bedrockConverseStream builds a converse-stream body: contentBlockDelta frames plus a trailing metadata usage frame.
func bedrockConverseStream(t *testing.T, deltas ...string) []byte {
t.Helper()
var body bytes.Buffer
for i, ev := range deltas {
eventType := "contentBlockDelta"
if i == len(deltas)-1 {
eventType = "metadata"
}
body.Write(awsFrame(t, eventType, []byte(ev)))
}
return body.Bytes()
}
proxy/internal/middleware/builtin/cost_meter/factory.go
+88
View File
@@ -0,0 +1,88 @@
package cost_meter
import (
"bytes"
"context"
"encoding/json"
"fmt"
"github.com/netbirdio/netbird/proxy/internal/llm/pricing"
"github.com/netbirdio/netbird/proxy/internal/middleware"
"github.com/netbirdio/netbird/proxy/internal/middleware/builtin"
)
// defaultPricingFilename is the basename probed inside the proxy data
// directory when no override is configured.
const defaultPricingFilename = "pricing.yaml"
// Config is the on-wire configuration for the middleware.
type Config struct {
// PricingPath optionally overrides the basename of the pricing
// file probed inside the proxy data directory. When empty the
// loader falls back to "pricing.yaml".
PricingPath string `json:"pricing_path"`
}
// Factory builds cost_meter instances from raw config bytes.
type Factory struct{}
// ID returns the registry identifier.
func (Factory) ID() string { return ID }
// New constructs a middleware instance. Empty, null, and {} configs
// are accepted; non-empty rawConfig that fails to unmarshal is
// rejected so misconfigurations surface at chain build time. The
// pricing loader is built once per instance and reused across
// invocations.
func (Factory) New(rawConfig []byte) (middleware.Middleware, error) {
cfg, err := decodeConfig(rawConfig)
if err != nil {
return nil, err
}
fctx := builtin.Context()
pricingPath := cfg.PricingPath
if pricingPath == "" {
pricingPath = defaultPricingFilename
}
loader, err := pricing.NewLoader(fctx.DataDir, pricingPath, ID, nil)
if err != nil {
return nil, fmt.Errorf("init pricing loader: %w", err)
}
cancel := startReloader(fctx.Context, loader)
return newMiddleware(loader, cancel), nil
}
// startReloader binds the loader's mtime-poll goroutine to a context
// derived from the proxy-lifetime context and returns its cancel func so
// the owning middleware can stop the goroutine on teardown. Returns nil
// when there's nothing to watch (nil context or defaults-only loader), in
// which case the middleware's Close is a no-op.
func startReloader(ctx context.Context, loader *pricing.Loader) context.CancelFunc {
if ctx == nil || !loader.WatchesFile() {
return nil
}
cctx, cancel := context.WithCancel(ctx)
go loader.Reload(cctx)
return cancel
}
// decodeConfig accepts empty, null, and {} configs, returning a
// zero-value Config. Non-empty payloads must parse cleanly.
func decodeConfig(rawConfig []byte) (Config, error) {
var cfg Config
if len(bytes.TrimSpace(rawConfig)) == 0 {
return cfg, nil
}
if err := json.Unmarshal(rawConfig, &cfg); err != nil {
return cfg, fmt.Errorf("decode config: %w", err)
}
return cfg, nil
}
func init() {
builtin.Register(Factory{})
}
proxy/internal/middleware/builtin/cost_meter/middleware.go
+218
View File
@@ -0,0 +1,218 @@
// Package cost_meter implements the SlotOnResponse middleware that
// converts token-usage metadata emitted by llm_response_parser into a
// per-request USD cost estimate. The middleware uses the shared pricing
// loader so operator pricing overrides apply to the chain.
package cost_meter
import (
"context"
"fmt"
"strconv"
"github.com/netbirdio/netbird/proxy/internal/llm/pricing"
"github.com/netbirdio/netbird/proxy/internal/middleware"
)
// ID is the registry identifier for this middleware.
const ID = "cost_meter"
// Version is the implementation version emitted via the spec merge.
const Version = "1.0.0"
// Skip reasons emitted under KeyCostSkipped. The set is closed; the
// dashboard surfaces these verbatim.
const (
skipMissingProvider = "missing_provider"
skipMissingModel = "missing_model"
skipMissingTokens = "missing_tokens"
//nolint:gosec // skip-reason label, not a credential
skipUnparseableTokens = "unparseable_tokens"
skipZeroTokens = "zero_tokens"
skipUnknownModel = "unknown_model"
)
var metadataKeys = []string{
middleware.KeyCostUSDInput,
middleware.KeyCostUSDCachedInput,
middleware.KeyCostUSDCacheCreation,
middleware.KeyCostUSDOutput,
middleware.KeyCostUSDTotal,
middleware.KeyCostUSDCache,
middleware.KeyCostSkipped,
}
// Middleware computes a per-response cost estimate from the token
// counts emitted upstream by llm_response_parser.
type Middleware struct {
loader *pricing.Loader
// cancel stops this instance's pricing-reload goroutine. Non-nil only
// when the loader watches an override file; Close calls it so a chain
// rebuild doesn't leak a poll goroutine per retired instance.
cancel context.CancelFunc
}
// newMiddleware constructs a Middleware bound to the given pricing loader.
// cancel may be nil (defaults-only loader with no reloader to stop).
func newMiddleware(loader *pricing.Loader, cancel context.CancelFunc) *Middleware {
return &Middleware{loader: loader, cancel: cancel}
}
// ID returns the registry identifier.
func (m *Middleware) ID() string { return ID }
// Version returns the implementation version.
func (m *Middleware) Version() string { return Version }
// Slot reports that the middleware runs after the upstream call.
func (m *Middleware) Slot() middleware.Slot { return middleware.SlotOnResponse }
// AcceptedContentTypes is empty: cost_meter never inspects bodies.
func (m *Middleware) AcceptedContentTypes() []string { return []string{} }
// MetadataKeys returns the closed allowlist of keys this middleware
// may emit.
func (m *Middleware) MetadataKeys() []string {
return append([]string(nil), metadataKeys...)
}
// MutationsSupported reports that this middleware never mutates the
// response.
func (m *Middleware) MutationsSupported() bool { return false }
// Close stops this instance's pricing-reload goroutine, if any. Called by
// the chain when a rebuild retires the instance, so the mtime-poll loop
// doesn't outlive the chain it belonged to. Safe to call on a nil receiver
// and on an instance with no reloader.
func (m *Middleware) Close() error {
if m != nil && m.cancel != nil {
m.cancel()
}
return nil
}
// Invoke reads provider, model, and token metadata, looks up pricing,
// and emits either KeyCostUSDTotal or KeyCostSkipped. The decision is
// always DecisionAllow; cost metering never denies or mutates.
func (m *Middleware) Invoke(_ context.Context, in *middleware.Input) (*middleware.Output, error) {
out := &middleware.Output{Decision: middleware.DecisionAllow}
if in == nil {
return out, nil
}
provider := lookupKV(in.Metadata, middleware.KeyLLMProvider)
if provider == "" {
out.Metadata = skip(skipMissingProvider)
return out, nil
}
model := lookupKV(in.Metadata, middleware.KeyLLMModel)
if model == "" {
out.Metadata = skip(skipMissingModel)
return out, nil
}
inRaw, hasIn := lookupKVOK(in.Metadata, middleware.KeyLLMInputTokens)
outRaw, hasOut := lookupKVOK(in.Metadata, middleware.KeyLLMOutputTokens)
if !hasIn || !hasOut {
out.Metadata = skip(skipMissingTokens)
return out, nil
}
inTokens, err := strconv.ParseInt(inRaw, 10, 64)
if err != nil || inTokens < 0 {
// Unparseable or negative tokens are not a runtime error: the
// upstream llm_response_parser emitted a non-numeric / invalid
// value, so we surface that as cost.skipped and continue with
// Allow rather than pricing a negative count.
out.Metadata = skip(skipUnparseableTokens)
return out, nil //nolint:nilerr // structured skip; not a runtime error
}
outTokens, err := strconv.ParseInt(outRaw, 10, 64)
if err != nil || outTokens < 0 {
out.Metadata = skip(skipUnparseableTokens)
return out, nil //nolint:nilerr // structured skip; not a runtime error
}
// Cache buckets are optional and silently zeroed on a missing /
// malformed value; they're a refinement on top of input cost,
// not a precondition. A buggy value falls back to 0, never aborts.
cachedTokens := parseOptionalInt64(in.Metadata, middleware.KeyLLMCachedInputTokens)
cacheCreationTokens := parseOptionalInt64(in.Metadata, middleware.KeyLLMCacheCreationTokens)
if inTokens == 0 && outTokens == 0 && cachedTokens == 0 && cacheCreationTokens == 0 {
out.Metadata = skip(skipZeroTokens)
return out, nil
}
table := m.loader.Get()
costs, ok := table.Costs(provider, model, inTokens, outTokens, cachedTokens, cacheCreationTokens)
if !ok {
out.Metadata = skip(skipUnknownModel)
return out, nil
}
// Per-bucket costs first: they're the base of the breakdown, and the two
// aggregates that follow are derived from exactly these four values.
out.Metadata = []middleware.KV{
{Key: middleware.KeyCostUSDInput, Value: usd(costs.InputUSD)},
{Key: middleware.KeyCostUSDCachedInput, Value: usd(costs.CachedInputUSD)},
{Key: middleware.KeyCostUSDCacheCreation, Value: usd(costs.CacheCreationUSD)},
{Key: middleware.KeyCostUSDOutput, Value: usd(costs.OutputUSD)},
{Key: middleware.KeyCostUSDTotal, Value: usd(costs.TotalUSD)},
{Key: middleware.KeyCostUSDCache, Value: usd(costs.CacheUSD)},
}
return out, nil
}
// usd renders a cost as the fixed-precision string every cost.usd_* key
// carries, so the per-bucket values and the aggregates round identically.
//
// 9 decimals, not 6: these values are summed downstream — per request, per
// session, and per usage bucket — so the rounding step is applied once per
// bucket per row and then accumulated. At 6 decimals a single row loses up to
// 2e-6 across its four buckets (enough to break a 1e-6 reconciliation against
// published rates), and a bucket smaller than half a microdollar quantises to
// zero outright: 16 cache-read tokens on a cheap model is 1.6e-9, so summing
// 10k such rows reports 0.02 instead of 0.016. Nano-dollar precision keeps the
// per-row error ~1000x below the smallest realistic bucket.
func usd(v float64) string { return fmt.Sprintf("%.9f", v) }
// skip returns a single-entry metadata slice carrying the given skip
// reason under KeyCostSkipped.
func skip(reason string) []middleware.KV {
return []middleware.KV{{Key: middleware.KeyCostSkipped, Value: reason}}
}
// lookupKV returns the value associated with key, or the empty string
// when the key is absent.
func lookupKV(kvs []middleware.KV, key string) string {
v, _ := lookupKVOK(kvs, key)
return v
}
// lookupKVOK returns the value associated with key plus a presence
// flag so callers can distinguish absent from empty.
func lookupKVOK(kvs []middleware.KV, key string) (string, bool) {
for _, kv := range kvs {
if kv.Key == key {
return kv.Value, true
}
}
return "", false
}
// parseOptionalInt64 reads a metadata value and decodes it as int64.
// Absent or unparseable values yield 0 — the caller treats absence as
// "no cached tokens" rather than an error, since cache buckets are a
// refinement, not a precondition.
func parseOptionalInt64(kvs []middleware.KV, key string) int64 {
raw, ok := lookupKVOK(kvs, key)
if !ok {
return 0
}
v, err := strconv.ParseInt(raw, 10, 64)
if err != nil || v < 0 {
return 0
}
return v
}
proxy/internal/middleware/builtin/cost_meter/middleware_test.go
+508
View File
@@ -0,0 +1,508 @@
package cost_meter
import (
"context"
"encoding/json"
"os"
"path/filepath"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/netbirdio/netbird/proxy/internal/middleware"
"github.com/netbirdio/netbird/proxy/internal/middleware/builtin"
)
const fixturePricing = `openai:
gpt-4o:
input_per_1k: 0.0025
output_per_1k: 0.01
gpt-4o-mini:
input_per_1k: 0.00015
output_per_1k: 0.0006
anthropic:
claude-sonnet-4-5:
input_per_1k: 0.003
output_per_1k: 0.015
`
// configureBuiltin points the package-level FactoryContext at a tmp
// directory containing the test pricing fixture. Returns the path so
// callers can override files later if needed.
func configureBuiltin(t *testing.T) string {
t.Helper()
dir := t.TempDir()
require.NoError(t, os.WriteFile(filepath.Join(dir, "pricing.yaml"), []byte(fixturePricing), 0o600), "write pricing fixture")
builtin.Configure(context.Background(), dir, nil, nil, nil)
return dir
}
func metaValue(t *testing.T, kvs []middleware.KV, key string) (string, bool) {
t.Helper()
for _, kv := range kvs {
if kv.Key == key {
return kv.Value, true
}
}
return "", false
}
func buildMiddleware(t *testing.T, raw []byte) middleware.Middleware {
t.Helper()
mw, err := Factory{}.New(raw)
require.NoError(t, err, "factory must accept the supplied config")
return mw
}
func TestMiddleware_StaticSurface(t *testing.T) {
configureBuiltin(t)
mw := buildMiddleware(t, nil)
assert.Equal(t, ID, mw.ID(), "ID must match the registered constant")
assert.Equal(t, Version, mw.Version(), "Version must match the constant")
assert.Equal(t, middleware.SlotOnResponse, mw.Slot(), "must run in the response slot")
assert.Empty(t, mw.AcceptedContentTypes(), "cost_meter does not inspect bodies")
assert.False(t, mw.MutationsSupported(), "cost_meter never mutates")
assert.NoError(t, mw.Close(), "Close on stateless middleware is a no-op")
keys := mw.MetadataKeys()
expected := []string{
middleware.KeyCostUSDInput,
middleware.KeyCostUSDCachedInput,
middleware.KeyCostUSDCacheCreation,
middleware.KeyCostUSDOutput,
middleware.KeyCostUSDTotal,
middleware.KeyCostUSDCache,
middleware.KeyCostSkipped,
}
assert.Equal(t, expected, keys, "metadata key allowlist must match the spec")
}
func TestFactory_AcceptsEmptyAndJSONConfig(t *testing.T) {
configureBuiltin(t)
cases := [][]byte{nil, {}, []byte("null"), []byte("{}"), []byte(" ")}
for _, raw := range cases {
mw, err := Factory{}.New(raw)
require.NoError(t, err, "empty/null/object config must be accepted")
require.NotNil(t, mw, "factory must return a middleware instance")
}
}
func TestFactory_RejectsMalformedConfig(t *testing.T) {
configureBuiltin(t)
mw, err := Factory{}.New([]byte("{not json"))
require.Error(t, err, "malformed config must surface at construction")
assert.Nil(t, mw, "no instance is returned on error")
}
func TestFactory_DefaultPricingPathLoadsFixture(t *testing.T) {
configureBuiltin(t)
mw := buildMiddleware(t, nil)
out, err := mw.Invoke(context.Background(), &middleware.Input{
Metadata: []middleware.KV{
{Key: middleware.KeyLLMProvider, Value: "openai"},
{Key: middleware.KeyLLMModel, Value: "gpt-4o-mini"},
{Key: middleware.KeyLLMInputTokens, Value: "1000"},
{Key: middleware.KeyLLMOutputTokens, Value: "1000"},
},
})
require.NoError(t, err)
require.Equal(t, middleware.DecisionAllow, out.Decision, "cost_meter always allows")
value, ok := metaValue(t, out.Metadata, middleware.KeyCostUSDTotal)
require.True(t, ok, "cost.usd_total must be emitted for known model")
assert.Equal(t, "0.000750000", value, "0.00015 + 0.0006 per 1k tokens, 9-decimal format")
}
func TestFactory_PricingPathOverride(t *testing.T) {
dir := t.TempDir()
require.NoError(t, os.WriteFile(filepath.Join(dir, "custom.yaml"), []byte(fixturePricing), 0o600), "write custom pricing")
builtin.Configure(context.Background(), dir, nil, nil, nil)
raw, err := json.Marshal(Config{PricingPath: "custom.yaml"})
require.NoError(t, err)
mw := buildMiddleware(t, raw)
out, err := mw.Invoke(context.Background(), &middleware.Input{
Metadata: []middleware.KV{
{Key: middleware.KeyLLMProvider, Value: "openai"},
{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
{Key: middleware.KeyLLMInputTokens, Value: "2000"},
{Key: middleware.KeyLLMOutputTokens, Value: "1000"},
},
})
require.NoError(t, err)
value, ok := metaValue(t, out.Metadata, middleware.KeyCostUSDTotal)
require.True(t, ok, "cost.usd_total must be emitted with custom pricing path")
assert.Equal(t, "0.015000000", value, "2*0.0025 + 1*0.01 = 0.015 with 9-decimal format")
}
func TestInvoke_ComputesCostForKnownModel(t *testing.T) {
configureBuiltin(t)
mw := buildMiddleware(t, nil)
out, err := mw.Invoke(context.Background(), &middleware.Input{
Metadata: []middleware.KV{
{Key: middleware.KeyLLMProvider, Value: "anthropic"},
{Key: middleware.KeyLLMModel, Value: "claude-sonnet-4-5"},
{Key: middleware.KeyLLMInputTokens, Value: "1000"},
{Key: middleware.KeyLLMOutputTokens, Value: "1000"},
},
})
require.NoError(t, err)
value, ok := metaValue(t, out.Metadata, middleware.KeyCostUSDTotal)
require.True(t, ok, "cost.usd_total must be emitted")
assert.Equal(t, "0.018000000", value, "0.003 + 0.015 = 0.018 with 9-decimal format")
_, skipped := metaValue(t, out.Metadata, middleware.KeyCostSkipped)
assert.False(t, skipped, "cost.skipped must not be set when cost is computed")
}
func TestInvoke_MissingProvider(t *testing.T) {
configureBuiltin(t)
mw := buildMiddleware(t, nil)
out, err := mw.Invoke(context.Background(), &middleware.Input{
Metadata: []middleware.KV{
{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
{Key: middleware.KeyLLMInputTokens, Value: "10"},
{Key: middleware.KeyLLMOutputTokens, Value: "10"},
},
})
require.NoError(t, err)
value, ok := metaValue(t, out.Metadata, middleware.KeyCostSkipped)
require.True(t, ok, "cost.skipped must be set when provider is missing")
assert.Equal(t, skipMissingProvider, value, "skip reason matches missing_provider")
}
func TestInvoke_MissingModel(t *testing.T) {
configureBuiltin(t)
mw := buildMiddleware(t, nil)
out, err := mw.Invoke(context.Background(), &middleware.Input{
Metadata: []middleware.KV{
{Key: middleware.KeyLLMProvider, Value: "openai"},
{Key: middleware.KeyLLMInputTokens, Value: "10"},
{Key: middleware.KeyLLMOutputTokens, Value: "10"},
},
})
require.NoError(t, err)
value, ok := metaValue(t, out.Metadata, middleware.KeyCostSkipped)
require.True(t, ok, "cost.skipped must be set when model is missing")
assert.Equal(t, skipMissingModel, value, "skip reason matches missing_model")
}
func TestInvoke_MissingTokens(t *testing.T) {
configureBuiltin(t)
mw := buildMiddleware(t, nil)
cases := []struct {
name string
md []middleware.KV
}{
{
name: "input only",
md: []middleware.KV{
{Key: middleware.KeyLLMProvider, Value: "openai"},
{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
{Key: middleware.KeyLLMInputTokens, Value: "10"},
},
},
{
name: "output only",
md: []middleware.KV{
{Key: middleware.KeyLLMProvider, Value: "openai"},
{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
{Key: middleware.KeyLLMOutputTokens, Value: "10"},
},
},
{
name: "neither",
md: []middleware.KV{
{Key: middleware.KeyLLMProvider, Value: "openai"},
{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
},
},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
out, err := mw.Invoke(context.Background(), &middleware.Input{Metadata: tc.md})
require.NoError(t, err)
value, ok := metaValue(t, out.Metadata, middleware.KeyCostSkipped)
require.True(t, ok, "cost.skipped must be set when token keys are missing")
assert.Equal(t, skipMissingTokens, value, "skip reason matches missing_tokens")
})
}
}
func TestInvoke_UnparseableTokens(t *testing.T) {
configureBuiltin(t)
mw := buildMiddleware(t, nil)
cases := []struct {
name string
in string
out string
}{
{name: "input non-numeric", in: "abc", out: "10"},
{name: "output non-numeric", in: "10", out: "xyz"},
{name: "both garbage", in: "??", out: "??"},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
out, err := mw.Invoke(context.Background(), &middleware.Input{
Metadata: []middleware.KV{
{Key: middleware.KeyLLMProvider, Value: "openai"},
{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
{Key: middleware.KeyLLMInputTokens, Value: tc.in},
{Key: middleware.KeyLLMOutputTokens, Value: tc.out},
},
})
require.NoError(t, err)
value, ok := metaValue(t, out.Metadata, middleware.KeyCostSkipped)
require.True(t, ok, "cost.skipped must be set on unparseable tokens")
assert.Equal(t, skipUnparseableTokens, value, "skip reason matches unparseable_tokens")
})
}
}
func TestInvoke_ZeroTokens(t *testing.T) {
configureBuiltin(t)
mw := buildMiddleware(t, nil)
out, err := mw.Invoke(context.Background(), &middleware.Input{
Metadata: []middleware.KV{
{Key: middleware.KeyLLMProvider, Value: "openai"},
{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
{Key: middleware.KeyLLMInputTokens, Value: "0"},
{Key: middleware.KeyLLMOutputTokens, Value: "0"},
},
})
require.NoError(t, err)
value, ok := metaValue(t, out.Metadata, middleware.KeyCostSkipped)
require.True(t, ok, "cost.skipped must be set when both token counts are zero")
assert.Equal(t, skipZeroTokens, value, "skip reason matches zero_tokens")
_, hasCost := metaValue(t, out.Metadata, middleware.KeyCostUSDTotal)
assert.False(t, hasCost, "cost.usd_total must not be emitted for zero tokens")
}
func TestInvoke_UnknownModel(t *testing.T) {
configureBuiltin(t)
mw := buildMiddleware(t, nil)
out, err := mw.Invoke(context.Background(), &middleware.Input{
Metadata: []middleware.KV{
{Key: middleware.KeyLLMProvider, Value: "openai"},
{Key: middleware.KeyLLMModel, Value: "fantasy-model-9000"},
{Key: middleware.KeyLLMInputTokens, Value: "10"},
{Key: middleware.KeyLLMOutputTokens, Value: "10"},
},
})
require.NoError(t, err)
value, ok := metaValue(t, out.Metadata, middleware.KeyCostSkipped)
require.True(t, ok, "cost.skipped must be set when pricing entry is absent")
assert.Equal(t, skipUnknownModel, value, "skip reason matches unknown_model")
}
func TestInvoke_NilInput(t *testing.T) {
configureBuiltin(t)
mw := buildMiddleware(t, nil)
out, err := mw.Invoke(context.Background(), nil)
require.NoError(t, err)
require.NotNil(t, out, "output must be returned even on nil input")
assert.Equal(t, middleware.DecisionAllow, out.Decision, "decision must be allow on nil input")
assert.Empty(t, out.Metadata, "no metadata must be emitted on nil input")
}
const fixturePricingWithCache = `openai:
gpt-4o:
input_per_1k: 0.0025
output_per_1k: 0.01
cached_input_per_1k: 0.00125
anthropic:
claude-sonnet-4-5:
input_per_1k: 0.003
output_per_1k: 0.015
cache_read_per_1k: 0.0003
cache_creation_per_1k: 0.00375
`
// configureBuiltinWithCacheRates points the package-level
// FactoryContext at a tmp directory containing pricing entries that
// include the cache rate fields.
func configureBuiltinWithCacheRates(t *testing.T) {
t.Helper()
dir := t.TempDir()
require.NoError(t, os.WriteFile(filepath.Join(dir, "pricing.yaml"), []byte(fixturePricingWithCache), 0o600), "write cache-aware pricing fixture")
builtin.Configure(context.Background(), dir, nil, nil, nil)
}
// TestInvoke_OpenAICachedSubsetDiscount proves the OpenAI shape end
// to end through the middleware: cached_input_tokens is treated as a
// SUBSET of input_tokens and discounted at the configured rate, not
// added on top.
func TestInvoke_OpenAICachedSubsetDiscount(t *testing.T) {
configureBuiltinWithCacheRates(t)
mw := buildMiddleware(t, nil)
out, err := mw.Invoke(context.Background(), &middleware.Input{
Metadata: []middleware.KV{
{Key: middleware.KeyLLMProvider, Value: "openai"},
{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
{Key: middleware.KeyLLMInputTokens, Value: "1000"},
{Key: middleware.KeyLLMOutputTokens, Value: "500"},
{Key: middleware.KeyLLMCachedInputTokens, Value: "750"},
},
})
require.NoError(t, err)
require.Equal(t, middleware.DecisionAllow, out.Decision)
value, ok := metaValue(t, out.Metadata, middleware.KeyCostUSDTotal)
require.True(t, ok, "cached subset path must produce a cost — never a skip")
// 250 non-cached at 0.0025/1k + 750 cached at 0.00125/1k + 500 output at 0.01/1k.
assert.Equal(t, "0.006562500", value,
"cached subset must be billed at the discount rate, non-cached at the full rate; never double-billed")
cache, ok := metaValue(t, out.Metadata, middleware.KeyCostUSDCache)
require.True(t, ok, "cost.usd_cache must be emitted alongside cost.usd_total")
// 750 cached at 0.00125/1k = 0.0009375.
assert.Equal(t, "0.000937500", cache, "cache cost is the discounted cost of the cached subset")
// Per-bucket breakdown. On OpenAI the cached subset is carved out of the
// input bucket, so input covers only the 250 non-cached tokens — the two
// must never double-count the same 750 tokens.
assertBucket(t, out.Metadata, middleware.KeyCostUSDInput, "0.000625000",
"input bucket bills only the non-cached remainder")
assertBucket(t, out.Metadata, middleware.KeyCostUSDCachedInput, "0.000937500",
"cached-input bucket bills the discounted subset")
assertBucket(t, out.Metadata, middleware.KeyCostUSDCacheCreation, "0.000000000",
"OpenAI has no cache-write bucket")
assertBucket(t, out.Metadata, middleware.KeyCostUSDOutput, "0.005000000",
"output bucket bills 500 tokens at 0.01/1k")
}
// TestInvoke_AnthropicCacheBucketsAdditive proves the Anthropic
// shape: cache_read and cache_creation are additive to input_tokens
// and each carries its own rate.
func TestInvoke_AnthropicCacheBucketsAdditive(t *testing.T) {
configureBuiltinWithCacheRates(t)
mw := buildMiddleware(t, nil)
out, err := mw.Invoke(context.Background(), &middleware.Input{
Metadata: []middleware.KV{
{Key: middleware.KeyLLMProvider, Value: "anthropic"},
{Key: middleware.KeyLLMModel, Value: "claude-sonnet-4-5"},
{Key: middleware.KeyLLMInputTokens, Value: "256"},
{Key: middleware.KeyLLMOutputTokens, Value: "200"},
{Key: middleware.KeyLLMCachedInputTokens, Value: "768"},
{Key: middleware.KeyLLMCacheCreationTokens, Value: "512"},
},
})
require.NoError(t, err)
require.Equal(t, middleware.DecisionAllow, out.Decision)
value, ok := metaValue(t, out.Metadata, middleware.KeyCostUSDTotal)
require.True(t, ok)
// 256 input * 0.003 + 768 cache_read * 0.0003 + 512 cache_creation * 0.00375 + 200 output * 0.015
// = 0.000768 + 0.0002304 + 0.00192 + 0.003 = 0.0059184.
assert.Equal(t, "0.005918400", value,
"each Anthropic input bucket must bill at its own rate — cache_read cheap, cache_creation expensive, regular input mid")
cache, ok := metaValue(t, out.Metadata, middleware.KeyCostUSDCache)
require.True(t, ok, "cost.usd_cache must be emitted alongside cost.usd_total")
// 768 cache_read * 0.0003 + 512 cache_creation * 0.00375 = 0.0021504.
assert.Equal(t, "0.002150400", cache, "cache cost sums the read and creation buckets")
// Per-bucket breakdown: four separately-billed buckets, each at its own rate.
assertBucket(t, out.Metadata, middleware.KeyCostUSDInput, "0.000768000",
"input bucket bills 256 tokens at 0.003/1k")
assertBucket(t, out.Metadata, middleware.KeyCostUSDCachedInput, "0.000230400",
"cache-read bucket bills 768 tokens at the cheap 0.0003/1k")
assertBucket(t, out.Metadata, middleware.KeyCostUSDCacheCreation, "0.001920000",
"cache-write bucket bills 512 tokens at the expensive 0.00375/1k")
assertBucket(t, out.Metadata, middleware.KeyCostUSDOutput, "0.003000000",
"output bucket bills 200 tokens at 0.015/1k")
}
// assertBucket asserts one per-bucket cost key carries the expected
// 6-decimal value.
func assertBucket(t *testing.T, md []middleware.KV, key, want, msg string) {
t.Helper()
got, ok := metaValue(t, md, key)
require.Truef(t, ok, "%s must be emitted", key)
assert.Equal(t, want, got, msg)
}
// TestInvoke_CachedTokensAbsentFallsBackToBaseFormula covers the
// "operator hasn't opted in" path: with no cached metadata keys
// emitted, the meter must produce exactly the same cost as before
// the feature landed. Critical so operators with the new binary but
// no YAML changes see no behavioural drift on OpenAI requests.
func TestInvoke_CachedTokensAbsentFallsBackToBaseFormula(t *testing.T) {
configureBuiltinWithCacheRates(t)
mw := buildMiddleware(t, nil)
out, err := mw.Invoke(context.Background(), &middleware.Input{
Metadata: []middleware.KV{
{Key: middleware.KeyLLMProvider, Value: "openai"},
{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
{Key: middleware.KeyLLMInputTokens, Value: "1000"},
{Key: middleware.KeyLLMOutputTokens, Value: "500"},
// No KeyLLMCachedInputTokens — the parser didn't see one.
},
})
require.NoError(t, err)
value, ok := metaValue(t, out.Metadata, middleware.KeyCostUSDTotal)
require.True(t, ok)
// 1000 input * 0.0025 + 500 output * 0.01 = 0.0025 + 0.005 = 0.0075
assert.Equal(t, "0.007500000", value, "no cached metadata = same cost as before the feature landed")
}
// TestInvoke_UnparseableCachedTokensSkippedSilently proves the
// optional-bucket contract: a malformed cached_input_tokens metadata
// value falls back to 0 (= no cached count) and continues with the
// regular formula. Cache buckets are a refinement, never a reason to
// abort cost computation.
func TestInvoke_UnparseableCachedTokensSkippedSilently(t *testing.T) {
configureBuiltinWithCacheRates(t)
mw := buildMiddleware(t, nil)
out, err := mw.Invoke(context.Background(), &middleware.Input{
Metadata: []middleware.KV{
{Key: middleware.KeyLLMProvider, Value: "openai"},
{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
{Key: middleware.KeyLLMInputTokens, Value: "1000"},
{Key: middleware.KeyLLMOutputTokens, Value: "500"},
{Key: middleware.KeyLLMCachedInputTokens, Value: "not-a-number"},
},
})
require.NoError(t, err)
value, ok := metaValue(t, out.Metadata, middleware.KeyCostUSDTotal)
require.True(t, ok, "garbage cache metadata must NOT switch the response from a cost to a skip — fall back to 0 cached")
assert.Equal(t, "0.007500000", value, "same as the no-cached-metadata path")
}
// TestMiddleware_CloseCancelsReloader proves Close stops the per-instance
// pricing-reload goroutine: a chain rebuild retires the old instance and
// calls Close, which must invoke the cancel func startReloader handed it so
// the mtime-poll loop doesn't outlive the chain.
func TestMiddleware_CloseCancelsReloader(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
m := newMiddleware(nil, cancel)
require.NoError(t, m.Close(), "Close must not error")
require.Error(t, ctx.Err(), "Close must cancel the reloader context so the poll goroutine exits")
}
// TestMiddleware_CloseNilSafe confirms Close is a no-op (no panic) for an
// instance with no reloader and for a nil receiver.
func TestMiddleware_CloseNilSafe(t *testing.T) {
require.NoError(t, newMiddleware(nil, nil).Close(), "no-reloader Close must be a no-op")
var m *Middleware
require.NoError(t, m.Close(), "nil-receiver Close must be safe")
}
proxy/internal/middleware/builtin/llm_guardrail/factory.go
+82
View File
@@ -0,0 +1,82 @@
package llm_guardrail
import (
"encoding/json"
"fmt"
"strings"
"github.com/netbirdio/netbird/proxy/internal/middleware"
"github.com/netbirdio/netbird/proxy/internal/middleware/builtin"
)
// Config is the JSON-decoded shape accepted by the factory. The
// runtime path consumes the normalised allowlist; raw config is not
// retained beyond construction.
type Config struct {
ModelAllowlist []string `json:"model_allowlist"`
PromptCapture PromptCapture `json:"prompt_capture"`
}
// PromptCapture toggles the optional prompt capture + redaction step
// that emits llm.request_prompt onto the metadata bag.
type PromptCapture struct {
Enabled bool `json:"enabled"`
RedactPii bool `json:"redact_pii"`
}
// Factory builds a configured llm_guardrail middleware instance.
type Factory struct{}
// ID returns the registry identifier matching the middleware ID.
func (Factory) ID() string { return ID }
// New decodes the raw JSON config and returns a ready Middleware. An
// empty / null / empty-object payload yields a zero-value Config.
func (Factory) New(rawConfig []byte) (middleware.Middleware, error) {
cfg := Config{}
if len(rawConfig) > 0 && !isEmptyJSON(rawConfig) {
if err := json.Unmarshal(rawConfig, &cfg); err != nil {
return nil, fmt.Errorf("decode config: %w", err)
}
}
return New(cfg), nil
}
// isEmptyJSON reports whether the payload is whitespace, null, or an
// empty object/array. The caller skips Unmarshal in that case so the
// zero-value Config flows through unchanged.
func isEmptyJSON(raw []byte) bool {
trimmed := strings.TrimSpace(string(raw))
switch trimmed {
case "", "null", "{}", "[]":
return true
}
return false
}
// normaliseConfig lowercases and trims allowlist entries so the runtime
// match is case-insensitive. Empty entries are dropped.
func normaliseConfig(cfg Config) Config {
if len(cfg.ModelAllowlist) == 0 {
return cfg
}
cleaned := make([]string, 0, len(cfg.ModelAllowlist))
for _, entry := range cfg.ModelAllowlist {
n := normaliseModel(entry)
if n == "" {
continue
}
cleaned = append(cleaned, n)
}
cfg.ModelAllowlist = cleaned
return cfg
}
// normaliseModel lowercases and trims a single model identifier.
func normaliseModel(model string) string {
return strings.ToLower(strings.TrimSpace(model))
}
func init() {
builtin.Register(Factory{})
}

Some files were not shown because too many files have changed in this diff Show More