Examples with DelegationTokenProvider io.pravega.client.security.auth.DelegationTokenProvider used on opensource projects

Example 1 with DelegationTokenProvider

use of io.pravega.client.security.auth.DelegationTokenProvider in project pravega by pravega.

the class SegmentMetadataClientImpl method sealSegmentAsync.

private CompletableFuture<SegmentSealed> sealSegmentAsync(Segment segment, DelegationTokenProvider tokenProvider) {
    log.trace("Sealing segment: {}", segment);
    RawClient connection = getConnection();
    long requestId = connection.getFlow().getNextSequenceNumber();
    return tokenProvider.retrieveToken().thenCompose(token -> connection.sendRequest(requestId, new SealSegment(requestId, segment.getScopedName(), token))).thenApply(r -> transformReply(r, SegmentSealed.class));
}

Example 2 with DelegationTokenProvider

use of io.pravega.client.security.auth.DelegationTokenProvider in project pravega by pravega.

the class BatchClientFactoryImpl method getStreamSegmentInfo.

private StreamSegmentsIterator getStreamSegmentInfo(final Stream stream, final StreamCut startStreamCut, final StreamCut endStreamCut) {
    log.debug("Start stream cut: {}, End stream cut: {}", startStreamCut, endStreamCut);
    StreamSegmentsInfoImpl.validateStreamCuts(startStreamCut, endStreamCut);
    StreamSegmentSuccessors segments = getAndHandleExceptions(controller.getSegments(startStreamCut, endStreamCut), RuntimeException::new);
    final SortedSet<Segment> segmentSet = new TreeSet<>(segments.getSegments());
    final DelegationTokenProvider tokenProvider = DelegationTokenProviderFactory.create(controller, stream.getScope(), stream.getStreamName(), AccessOperation.READ);
    log.debug("List of Segments between the start and end stream cuts : {}", segmentSet);
    val futures = segmentSet.stream().map(s -> getSegmentRange(s, startStreamCut, endStreamCut, tokenProvider)).collect(Collectors.toList());
    List<SegmentRange> results = Futures.getThrowingException(Futures.allOfWithResults(futures));
    return StreamSegmentsInfoImpl.builder().segmentRangeIterator(results.iterator()).startStreamCut(startStreamCut).endStreamCut(endStreamCut).build();
}

Example 3 with DelegationTokenProvider

use of io.pravega.client.security.auth.DelegationTokenProvider in project pravega by pravega.

the class AsyncSegmentInputStreamTest method testProcessingFailure.

@Test
public void testProcessingFailure() throws ConnectionFailedException {
    Segment segment = new Segment("scope", "testRetry", 4);
    PravegaNodeUri endpoint = new PravegaNodeUri("localhost", SERVICE_PORT);
    @Cleanup MockConnectionFactoryImpl connectionFactory = new MockConnectionFactoryImpl();
    @Cleanup MockController controller = new MockController(endpoint.getEndpoint(), endpoint.getPort(), connectionFactory, true);
    DelegationTokenProvider tokenProvider = mock(DelegationTokenProvider.class);
    // return empty token
    when(tokenProvider.retrieveToken()).thenReturn(CompletableFuture.completedFuture(""));
    Semaphore dataAvailable = new Semaphore(0);
    @Cleanup AsyncSegmentInputStreamImpl in = new AsyncSegmentInputStreamImpl(controller, connectionFactory, segment, tokenProvider, dataAvailable);
    ClientConnection c = mock(ClientConnection.class);
    InOrder inOrder = Mockito.inOrder(c);
    connectionFactory.provideConnection(endpoint, c);
    WireCommands.SegmentRead segmentRead = new WireCommands.SegmentRead(segment.getScopedName(), 1234, false, false, Unpooled.EMPTY_BUFFER, in.getRequestId());
    Mockito.doAnswer(new Answer<Void>() {

        @Override
        public Void answer(InvocationOnMock invocation) throws Throwable {
            connectionFactory.getProcessor(endpoint).processingFailure(new ConnectionFailedException("Custom error"));
            return null;
        }
    }).doAnswer(new Answer<Void>() {

        @Override
        public Void answer(InvocationOnMock invocation) throws Throwable {
            connectionFactory.getProcessor(endpoint).process(segmentRead);
            return null;
        }
    }).when(c).send(any(ReadSegment.class));
    assertEquals(0, dataAvailable.availablePermits());
    CompletableFuture<SegmentRead> readFuture = in.read(1234, 5678);
    assertEquals(segmentRead, readFuture.join());
    assertTrue(Futures.isSuccessful(readFuture));
    assertEquals(1, dataAvailable.availablePermits());
    inOrder.verify(c).send(eq(new WireCommands.ReadSegment(segment.getScopedName(), 1234, 5678, "", in.getRequestId())));
    inOrder.verify(c).close();
    inOrder.verify(c).send(eq(new WireCommands.ReadSegment(segment.getScopedName(), 1234, 5678, "", in.getRequestId())));
    verifyNoMoreInteractions(c);
    // ensure retrieve Token is invoked for every retry.
    verify(tokenProvider, times(2)).retrieveToken();
}

Example 4 with DelegationTokenProvider

use of io.pravega.client.security.auth.DelegationTokenProvider in project pravega by pravega.

the class TransactionalEventStreamWriterImpl method beginTxn.

@Override
public Transaction<Type> beginTxn() {
    TxnSegments txnSegments = Futures.getThrowingException(controller.createTransaction(stream, config.getTransactionTimeoutTime()));
    log.info("Transaction {} created", txnSegments.getTxnId());
    UUID txnId = txnSegments.getTxnId();
    Map<Segment, SegmentTransaction<Type>> transactions = new HashMap<>();
    DelegationTokenProvider tokenProvider = null;
    for (Segment s : txnSegments.getStreamSegments().getSegments()) {
        if (tokenProvider == null) {
            tokenProvider = DelegationTokenProviderFactory.create(controller, s, AccessOperation.WRITE);
        }
        SegmentOutputStream out = outputStreamFactory.createOutputStreamForTransaction(s, txnId, config, tokenProvider);
        SegmentTransactionImpl<Type> impl = new SegmentTransactionImpl<>(txnId, out, serializer);
        transactions.put(s, impl);
    }
    pinger.startPing(txnId);
    return new TransactionImpl<Type>(writerId, txnId, transactions, txnSegments.getStreamSegments(), controller, stream, pinger);
}

Example 5 with DelegationTokenProvider

use of io.pravega.client.security.auth.DelegationTokenProvider in project pravega by pravega.

the class TransactionalEventStreamWriterImpl method getTxn.

@Override
public Transaction<Type> getTxn(UUID txId) {
    // check if the transaction is open.
    Status status = Futures.getThrowingException(controller.checkTransactionStatus(stream, txId));
    if (status != Status.OPEN) {
        return new TransactionImpl<>(writerId, txId, controller, stream);
    }
    // get the segments corresponding to the transaction.
    StreamSegments segments = Futures.getThrowingException(controller.getEpochSegments(stream.getScope(), stream.getStreamName(), getEpoch(txId)));
    assert segments != null : "Epoch segments returned is null";
    Preconditions.checkState(segments.getSegments().size() > 0, "There should be at least 1 epoch segment");
    // Create OutputStream for every segment.
    Map<Segment, SegmentTransaction<Type>> transactions = new HashMap<>();
    DelegationTokenProvider tokenProvider = null;
    for (Segment s : segments.getSegments()) {
        if (tokenProvider == null) {
            tokenProvider = DelegationTokenProviderFactory.create(controller, s, AccessOperation.WRITE);
        }
        SegmentOutputStream out = outputStreamFactory.createOutputStreamForTransaction(s, txId, config, tokenProvider);
        SegmentTransactionImpl<Type> impl = new SegmentTransactionImpl<>(txId, out, serializer);
        transactions.put(s, impl);
    }
    pinger.startPing(txId);
    return new TransactionImpl<Type>(writerId, txId, transactions, segments, controller, stream, pinger);
}