Skip to content

Commit d3b1d09

Browse files
AntonLiashenkaAntonLiashenka
authored
fix-issues (#81)
* fix resolve domain * fix host header resolution and extraction
1 parent 65c7aa1 commit d3b1d09

6 files changed

Lines changed: 188 additions & 21 deletions

File tree

‎.gitignore‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,3 +48,4 @@ log
4848

4949
.junie
5050

51+
/CLAUDE.md

‎cli/src/main/java/tech/amak/portbuddy/cli/tunnel/HttpTunnelClient.java‎

Lines changed: 47 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -218,6 +218,45 @@ public void close() {
218218
}
219219
}
220220

221+
/**
222+
* Resolves the Host header to present to the local target. Prefers the server-supplied
223+
* {@code X-Forwarded-Host} (which reflects whichever public domain the visitor actually
224+
* used — the portbuddy.dev subdomain or an attached custom domain), falling back to the
225+
* tunnel's original public URL only if the server didn't send one (older server versions).
226+
*/
227+
private String resolveLocalHostHeader(final String forwardedHost) {
228+
if (forwardedHost != null && !forwardedHost.isBlank()) {
229+
final var commaIdx = forwardedHost.indexOf(',');
230+
return (commaIdx > 0 ? forwardedHost.substring(0, commaIdx) : forwardedHost).trim();
231+
}
232+
return URI.create(publicBaseUrl).getHost();
233+
}
234+
235+
private static String extractHeaderValue(final Map<String, List<String>> headers, final String name) {
236+
if (headers == null) {
237+
return null;
238+
}
239+
for (final var entry : headers.entrySet()) {
240+
if (entry.getKey() != null && entry.getKey().equalsIgnoreCase(name)
241+
&& entry.getValue() != null && !entry.getValue().isEmpty()) {
242+
return entry.getValue().getFirst();
243+
}
244+
}
245+
return null;
246+
}
247+
248+
private static String extractSingleHeaderValue(final Map<String, String> headers, final String name) {
249+
if (headers == null) {
250+
return null;
251+
}
252+
for (final var entry : headers.entrySet()) {
253+
if (entry.getKey() != null && entry.getKey().equalsIgnoreCase(name) && entry.getValue() != null) {
254+
return entry.getValue();
255+
}
256+
}
257+
return null;
258+
}
259+
221260
private String toWebSocketUrl(final String base, final String path) {
222261
final var uri = URI.create(base);
223262
var scheme = uri.getScheme();
@@ -357,9 +396,10 @@ private void handleWsFromServer(final WsTunnelMessage message) {
357396
url += "?" + message.getQuery();
358397
}
359398
final var builder = new Request.Builder().url(url);
360-
final var publicHost = URI.create(publicBaseUrl).getHost();
361-
if (publicHost != null) {
362-
builder.header("Host", publicHost);
399+
final var hostHeader = resolveLocalHostHeader(extractSingleHeaderValue(message.getHeaders(),
400+
"X-Forwarded-Host"));
401+
if (hostHeader != null) {
402+
builder.header("Host", hostHeader);
363403
}
364404
if (message.getHeaders() != null) {
365405
for (final var entry : message.getHeaders().entrySet()) {
@@ -474,9 +514,10 @@ private HttpTunnelMessage handleRequest(final HttpTunnelMessage requestMessage)
474514
.url(url)
475515
.method(method, buildBody(method, requestMessage.getBodyB64(), requestMessage.getBodyContentType()));
476516

477-
final var publicHost = URI.create(publicBaseUrl).getHost();
478-
if (publicHost != null) {
479-
targetRequest.header("Host", publicHost);
517+
final var hostHeader = resolveLocalHostHeader(extractHeaderValue(requestMessage.getHeaders(),
518+
"X-Forwarded-Host"));
519+
if (hostHeader != null) {
520+
targetRequest.header("Host", hostHeader);
480521
}
481522

482523
if (requestMessage.getHeaders() != null) {

‎server/src/main/java/tech/amak/portbuddy/server/db/repo/AccountRepository.java‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,16 +20,27 @@
2020
import java.util.UUID;
2121

2222
import org.springframework.data.jpa.repository.JpaRepository;
23+
import org.springframework.data.jpa.repository.Lock;
2324
import org.springframework.data.jpa.repository.Query;
2425
import org.springframework.data.repository.query.Param;
2526

27+
import jakarta.persistence.LockModeType;
2628
import tech.amak.portbuddy.server.db.entity.AccountEntity;
2729
import tech.amak.portbuddy.server.web.admin.dto.AdminAccountRow;
2830
import tech.amak.portbuddy.server.web.admin.dto.AdminStatsRow;
2931

3032
public interface AccountRepository extends JpaRepository<AccountEntity, UUID> {
3133
Optional<AccountEntity> findByStripeCustomerId(String stripeCustomerId);
3234

35+
/**
36+
* Locks the account row for the duration of the current transaction. Used to serialize
37+
* concurrent domain-resolution/tunnel-creation requests for the same account so two
38+
* simultaneous expose calls can't both see the same domain as available.
39+
*/
40+
@Lock(LockModeType.PESSIMISTIC_WRITE)
41+
@Query("SELECT a FROM AccountEntity a WHERE a.id = :id")
42+
Optional<AccountEntity> findByIdForUpdate(@Param("id") UUID id);
43+
3344
@Query("SELECT a FROM AccountEntity a WHERE a.subscriptionStatus <> 'active' AND a.updatedAt < :cutoff")
3445
List<AccountEntity> findBySubscriptionStatusNotActiveAndUpdatedAtBefore(@Param("cutoff") OffsetDateTime cutoff);
3546

‎server/src/main/java/tech/amak/portbuddy/server/db/repo/TunnelRepository.java‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -36,8 +36,6 @@
3636
@Repository
3737
public interface TunnelRepository extends JpaRepository<TunnelEntity, UUID> {
3838

39-
boolean existsByDomainAndStatus(DomainEntity domain, TunnelStatus status);
40-
4139
boolean existsByDomainAndStatusNot(DomainEntity domain, TunnelStatus status);
4240

4341
Optional<TunnelEntity> findFirstByAccountIdAndLocalHostAndLocalPortAndDomainIsNotNullOrderByCreatedAtDesc(

‎server/src/main/java/tech/amak/portbuddy/server/service/DomainService.java‎

Lines changed: 36 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -20,16 +20,19 @@
2020
import java.util.Objects;
2121
import java.util.Optional;
2222
import java.util.UUID;
23+
import java.util.stream.Collectors;
2324
import javax.naming.Context;
2425
import javax.naming.NamingException;
2526
import javax.naming.directory.Attribute;
2627
import javax.naming.directory.Attributes;
2728
import javax.naming.directory.DirContext;
2829
import javax.naming.directory.InitialDirContext;
2930

31+
import org.springframework.http.HttpStatus;
3032
import org.springframework.security.crypto.password.PasswordEncoder;
3133
import org.springframework.stereotype.Service;
3234
import org.springframework.transaction.annotation.Transactional;
35+
import org.springframework.web.server.ResponseStatusException;
3336

3437
import lombok.RequiredArgsConstructor;
3538
import lombok.extern.slf4j.Slf4j;
@@ -38,6 +41,7 @@
3841
import tech.amak.portbuddy.server.db.entity.AccountEntity;
3942
import tech.amak.portbuddy.server.db.entity.DomainEntity;
4043
import tech.amak.portbuddy.server.db.entity.TunnelStatus;
44+
import tech.amak.portbuddy.server.db.repo.AccountRepository;
4145
import tech.amak.portbuddy.server.db.repo.DomainRepository;
4246
import tech.amak.portbuddy.server.db.repo.TunnelRepository;
4347
import tech.amak.portbuddy.server.db.repo.UserRepository;
@@ -55,6 +59,7 @@ public class DomainService {
5559
private final PasswordEncoder passwordEncoder;
5660
private final SslServiceClient sslServiceClient;
5761
private final UserRepository userRepository;
62+
private final AccountRepository accountRepository;
5863

5964
private final SecureRandom random = new SecureRandom();
6065

@@ -345,6 +350,11 @@ public DomainEntity resolveDomain(final AccountEntity account,
345350
final String requestedDomain,
346351
final String localHost,
347352
final Integer localPort) {
353+
// Lock the account row for the rest of this transaction so two concurrent expose
354+
// requests for the same account can't both read the same domain as available and
355+
// double-assign it before either tunnel commits.
356+
accountRepository.findByIdForUpdate(account.getId());
357+
348358
if (requestedDomain != null && !requestedDomain.isBlank()) {
349359
// User requested specific domain
350360
final var normalizedDomain = requestedDomain.toLowerCase();
@@ -356,19 +366,24 @@ public DomainEntity resolveDomain(final AccountEntity account,
356366

357367
final var finalSubdomain = targetSubdomain;
358368

359-
return domainRepository.findByAccountAndSubdomain(account, finalSubdomain)
360-
.filter(domain -> !isTunnelConnected(domain))
361-
.orElseThrow(() -> new RuntimeException("Domain not found or unavailable: " + requestedDomain));
369+
final var domain = domainRepository.findByAccountAndSubdomain(account, finalSubdomain)
370+
.filter(d -> !isTunnelActive(d))
371+
.orElseThrow(() -> new ResponseStatusException(HttpStatus.CONFLICT,
372+
"Domain not found or unavailable: " + requestedDomain));
373+
log.info("Resolved domain {} for account {} (explicit request)", finalSubdomain, account.getId());
374+
return domain;
362375
}
363376

364-
// No specific domain requested
365-
// Filter out CONNECTED domains
377+
// No specific domain requested.
378+
// Filter out domains with any non-closed tunnel (CONNECTED or still-PENDING), since a
379+
// PENDING tunnel's owner may finish connecting at any moment.
366380
final var availableDomains = domainRepository.findAllByAccount(account).stream()
367-
.filter(domain -> !isTunnelConnected(domain))
381+
.filter(domain -> !isTunnelActive(domain))
368382
.toList();
369383

370384
if (availableDomains.isEmpty()) {
371-
throw new RuntimeException("No available domains found. Please add a new domain at https://portbuddy.dev/app/domains");
385+
throw new ResponseStatusException(HttpStatus.CONFLICT,
386+
"No available domains found. Please add a new domain at https://portbuddy.dev/app/domains");
372387
}
373388

374389
// Affinity check: Find the last used subdomain for this resource
@@ -380,16 +395,25 @@ public DomainEntity resolveDomain(final AccountEntity account,
380395
.filter(domain -> Objects.equals(domain.getId(), lastDomain.getId()))
381396
.findFirst();
382397
if (matched.isPresent()) {
398+
log.info("Resolved domain {} for account {} (affinity match on {}:{})",
399+
matched.get().getSubdomain(), account.getId(), localHost, localPort);
383400
return matched.get();
384401
}
385402
}
386403

387-
// Pick any
388-
return availableDomains.getFirst();
389-
}
404+
if (availableDomains.size() > 1) {
405+
final var candidates = availableDomains.stream()
406+
.map(DomainEntity::getSubdomain)
407+
.collect(Collectors.joining(", "));
408+
throw new ResponseStatusException(HttpStatus.CONFLICT,
409+
"Multiple domains available for this account; specify --domain explicitly. Candidates: "
410+
+ candidates);
411+
}
390412

391-
private boolean isTunnelConnected(final DomainEntity domain) {
392-
return tunnelRepository.existsByDomainAndStatus(domain, TunnelStatus.CONNECTED);
413+
final var onlyDomain = availableDomains.getFirst();
414+
log.info("Resolved domain {} for account {} (sole available domain)",
415+
onlyDomain.getSubdomain(), account.getId());
416+
return onlyDomain;
393417
}
394418

395419
private boolean isTunnelActive(final DomainEntity domain) {

‎server/src/test/java/tech/amak/portbuddy/server/service/DomainServiceTest.java‎

Lines changed: 93 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,10 +22,13 @@
2222
import static org.junit.jupiter.api.Assertions.assertTrue;
2323
import static org.mockito.ArgumentMatchers.any;
2424
import static org.mockito.ArgumentMatchers.anyString;
25+
import static org.mockito.ArgumentMatchers.eq;
26+
import static org.mockito.Mockito.lenient;
2527
import static org.mockito.Mockito.times;
2628
import static org.mockito.Mockito.verify;
2729
import static org.mockito.Mockito.when;
2830

31+
import java.util.List;
2932
import java.util.Optional;
3033
import java.util.UUID;
3134

@@ -36,12 +39,15 @@
3639
import org.mockito.junit.jupiter.MockitoExtension;
3740
import org.springframework.security.crypto.password.PasswordEncoder;
3841
import org.springframework.util.unit.DataSize;
42+
import org.springframework.web.server.ResponseStatusException;
3943

4044
import tech.amak.portbuddy.server.client.SslServiceClient;
4145
import tech.amak.portbuddy.server.config.AppProperties;
4246
import tech.amak.portbuddy.server.db.entity.AccountEntity;
4347
import tech.amak.portbuddy.server.db.entity.DomainEntity;
48+
import tech.amak.portbuddy.server.db.entity.TunnelEntity;
4449
import tech.amak.portbuddy.server.db.entity.TunnelStatus;
50+
import tech.amak.portbuddy.server.db.repo.AccountRepository;
4551
import tech.amak.portbuddy.server.db.repo.DomainRepository;
4652
import tech.amak.portbuddy.server.db.repo.TunnelRepository;
4753
import tech.amak.portbuddy.server.db.repo.UserRepository;
@@ -61,6 +67,8 @@ class DomainServiceTest {
6167
private PasswordEncoder passwordEncoder;
6268
@Mock
6369
private SslServiceClient sslServiceClient;
70+
@Mock
71+
private AccountRepository accountRepository;
6472

6573
private DomainService domainService;
6674
private AccountEntity account;
@@ -93,9 +101,93 @@ void setUp() {
93101
appProps,
94102
passwordEncoder,
95103
sslServiceClient,
96-
userRepository);
104+
userRepository,
105+
accountRepository);
97106
account = new AccountEntity();
98107
account.setId(UUID.randomUUID());
108+
lenient().when(accountRepository.findByIdForUpdate(account.getId())).thenReturn(Optional.of(account));
109+
}
110+
111+
private static DomainEntity domainWithSubdomain(final String subdomain) {
112+
final var domain = new DomainEntity();
113+
domain.setId(UUID.randomUUID());
114+
domain.setSubdomain(subdomain);
115+
return domain;
116+
}
117+
118+
@Test
119+
void resolveDomain_NoRequest_SoleAvailableDomain_ReturnsIt() {
120+
final var domain = domainWithSubdomain("dieta18-admin");
121+
when(domainRepository.findAllByAccount(account)).thenReturn(List.of(domain));
122+
when(tunnelRepository.existsByDomainAndStatusNot(domain, TunnelStatus.CLOSED)).thenReturn(false);
123+
when(tunnelRepository.findUsedTunnel(account.getId(), "localhost", 80)).thenReturn(Optional.empty());
124+
125+
final var resolved = domainService.resolveDomain(account, null, "localhost", 80);
126+
127+
assertEquals(domain, resolved);
128+
verify(accountRepository).findByIdForUpdate(account.getId());
129+
}
130+
131+
@Test
132+
void resolveDomain_NoRequest_AffinityMatch_ReturnsMatchedDomain() {
133+
final var glubokoe = domainWithSubdomain("glubokoe-admin");
134+
final var dieta18 = domainWithSubdomain("dieta18-admin");
135+
when(domainRepository.findAllByAccount(account)).thenReturn(List.of(glubokoe, dieta18));
136+
when(tunnelRepository.existsByDomainAndStatusNot(any(DomainEntity.class), eq(TunnelStatus.CLOSED)))
137+
.thenReturn(false);
138+
139+
final var lastTunnel = new TunnelEntity();
140+
lastTunnel.setDomain(glubokoe);
141+
when(tunnelRepository.findUsedTunnel(account.getId(), "localhost", 8080)).thenReturn(Optional.of(lastTunnel));
142+
143+
final var resolved = domainService.resolveDomain(account, null, "localhost", 8080);
144+
145+
assertEquals(glubokoe, resolved);
146+
}
147+
148+
@Test
149+
void resolveDomain_NoRequest_MultipleAvailableNoAffinity_ThrowsConflictInsteadOfGuessing() {
150+
final var glubokoe = domainWithSubdomain("glubokoe-admin");
151+
final var dieta18 = domainWithSubdomain("dieta18-admin");
152+
when(domainRepository.findAllByAccount(account)).thenReturn(List.of(glubokoe, dieta18));
153+
when(tunnelRepository.existsByDomainAndStatusNot(any(DomainEntity.class), eq(TunnelStatus.CLOSED)))
154+
.thenReturn(false);
155+
when(tunnelRepository.findUsedTunnel(account.getId(), "localhost", 80)).thenReturn(Optional.empty());
156+
157+
final var ex = assertThrows(ResponseStatusException.class,
158+
() -> domainService.resolveDomain(account, null, "localhost", 80));
159+
assertEquals(409, ex.getStatusCode().value());
160+
}
161+
162+
@Test
163+
void resolveDomain_NoRequest_NoDomainsAvailable_ThrowsConflict() {
164+
when(domainRepository.findAllByAccount(account)).thenReturn(List.of());
165+
166+
assertThrows(ResponseStatusException.class,
167+
() -> domainService.resolveDomain(account, null, "localhost", 80));
168+
}
169+
170+
@Test
171+
void resolveDomain_ExplicitRequest_PendingTunnelBlocksDomain() {
172+
final var domain = domainWithSubdomain("dieta18-admin");
173+
when(domainRepository.findByAccountAndSubdomain(account, "dieta18-admin")).thenReturn(Optional.of(domain));
174+
// Domain has a non-CLOSED (e.g. PENDING) tunnel, so it must be treated as unavailable
175+
// even though it is not yet CONNECTED.
176+
when(tunnelRepository.existsByDomainAndStatusNot(domain, TunnelStatus.CLOSED)).thenReturn(true);
177+
178+
assertThrows(ResponseStatusException.class,
179+
() -> domainService.resolveDomain(account, "dieta18-admin", "localhost", 80));
180+
}
181+
182+
@Test
183+
void resolveDomain_ExplicitRequest_Available_ReturnsIt() {
184+
final var domain = domainWithSubdomain("dieta18-admin");
185+
when(domainRepository.findByAccountAndSubdomain(account, "dieta18-admin")).thenReturn(Optional.of(domain));
186+
when(tunnelRepository.existsByDomainAndStatusNot(domain, TunnelStatus.CLOSED)).thenReturn(false);
187+
188+
final var resolved = domainService.resolveDomain(account, "dieta18-admin", "localhost", 80);
189+
190+
assertEquals(domain, resolved);
99191
}
100192

101193
@Test

0 commit comments

Comments
 (0)